Spatio-Temporal Pruning for Compressed Spiking Large Language Models

TL;DR

This paper introduces a novel spatio-temporal pruning framework for Spiking Large Language Models, significantly reducing computational complexity and inference latency while maintaining high performance, thus enabling energy-efficient NLP applications on edge devices.

Contribution

It presents the first combined approach of spatial and temporal pruning, quantization, and distillation for Spiking LLMs, enhancing their efficiency and practicality.

Findings

Reduced active neurons and attention heads with spatial pruning.

Minimized inference latency through dynamic temporal adjustment.

Demonstrated improved efficiency on GLUE benchmark with SpikingBERT.

Abstract

Large Language Models (LLMs) present significant challenges for deployment in energy-constrained environments due to their large model sizes and high inference latency. Spiking Neural Networks (SNNs), inspired by the sparse event-driven neural processing and energy-efficient information transmission in the brain, offer a promising alternative for achieving low-power computing. Integrating the event-driven efficiency of spiking neurons with the advanced capabilities of LLMs represents a promising direction for power-efficient LLMs. This work specifically delves into the design of compressed spiking LLMs. Here, we revisit spatial and temporal pruning from the perspective of SNNs and propose a novel spatio-temporal pruning framework for Spiking LLMs to optimize computational efficiency while preserving high performance. Our spatial pruning technique reduces the number of active neurons and attention heads, effectively lowering the computational complexity of the model. Meanwhile, temporal pruning minimizes inference latency by dynamically adjusting the number of timesteps required for different layers. By combining these approaches with other compression techniques, we present the first work in the domain of Spiking LLMs to jointly explore spatial pruning, temporal pruning, extreme quantization and knowledge distillation strategies. Extensive experimental evaluation of our proposed framework for SpikingBERT on the large-scale GLUE benchmark demonstrates the efficacy of our approach in terms of computational operations and inference latency. Our approach offers a compelling solution for real-time, low-power natural language processing applications, making Spiking LLMs more practical for deployment on edge devices and in power-constrained settings.

Tables

Table 1. TABLE I: Hyperparameters Used During Pruning of SpikingBERT Models for Various Datasets

Dataset	MNLI	QQP	QNLI	SST-2	MRPC
$T_{conv}$	140	80	110	85	110
Initial $V_{th}$	0.85	1.0	0.9	1.0	0.9
Computational Cost Penalty $\lambda$	1e-12	2e-12	2e-12	5e-12	5e-12
PCA # of Components	0.99999	0.99999	0.99999	0.99999	0.99999
Training Batch Size	32	80	16	16	32
Testing Batch Size	128	128	128	32	128
Learning Rate	5e-8	1e-7	1e-7	5e-10	1e-8
ACs Constraint	0.6	0.6	0.6	0.6	0.6
PCA Scaling Factor $b$	1.02	1.02	1.02	1.02	1.02
Activity Loss Scaling Factor $\eta$	0.002	0.001	0.0015	0.001	0.0005

Table 2. TABLE II: Performance Comparison of Various Pruning Methods on 1-Bit SpikingBERT Model Across Different Tasks

Dataset	MNLI		QQP		QNLI		SST-2		MRPC
	Acc	ACs	Acc	ACs	Acc	ACs	Acc	ACs	Acc	ACs
	(%)	(%)	(%)	(%)	(%)	(%)	(%)	(%)	(%)	(%)
Baseline	75.54	100	83.99	100	80.85	100	86.58	100	75.71	100
SPATIAL PRUNING
Post-Training (FIM)	72.03	60	82.17	60	80.27	60	85.32	60	74.43	60
Post-Training (FIM + ASR of Neurons)	72.77	60	82.63	60	80.67	60	85.67	60	74.67	60
Post-Training (FIM + ASR of Attn)	72.5	60	82.53	60	80.52	60	85.44	60	74.84	60
Post-Training (FIM + ASR)	73.17	60	82.84	60	80.82	60	85.89	60	75.01	60
Post-Training (FIM + ASR) + Retraining	74.97	54.17	84.24	57.95	81.88	39.96	86.47	60	76.12	33.18
TEMPORAL PRUNING
Post-Training (PCA)	75.26	56.05	82.63	64.22	77.16	62.87	86.47	61.68	74.61	57.96
Post-Training (PCA) + Retraining	75.33	56.05	84.01	68.66	81.93	62.87	86.70	47.14	75.71	59.72
COMBINED SPATIAL AND TEMPORAL PRUNING
Two-Stage Pruning	75.13	41.67	84.10	41.82	82.28	28.26	87.27	41.53	76.12	18.19

Table 3. TABLE III: Performance Improvement Due to Pruning and Adaptive Thresholding on 1-Bit SpikingBERT Models Across Different Tasks

Dataset	Baseline model			Pruned Model			Adaptive Threshold
	Accuracy	Normalized #C	Latency	Accuracy	Normalized #C	Latency	Accuracy	Normalized #C	Latency
MNLI	75.54	0.219	140	75.13	0.080	99.6	73.71	0.022	28.2
QQP	83.99	0.266	80	84.10	0.103	55	82.77	0.027	15.20
QNLI	80.85	0.239	110	82.28	0.064	69.56	79.97	0.017	19.40
SST-2	86.58	0.241	85	87.27	0.100	53.08	86.58	0.023	13.04
MRPC	75.71	0.272	110	76.12	0.043	54	74.67	0.013	17.48

Table 4. TABLE IV: Comparison of Post-Training Pruning Strategies on SST-2 Dataset

Pruning Strategy	Spatial	Temporal	Retrain	Accuracy
No Pruning (Baseline)	✗	✗	0	86.58
Spatial Only	✓	✗	1	86.47
Temporal Only	✗	✓	1	86.70
Pruning + Joint Retraining	✓	✓	1	85.55
Sequential (Ours)	✓	✓	2	87.27

Table 5. TABLE V: Accuracy Comparison on SST-2 Dataset Between Baseline and Pruned Models Under Different Threshold Types and Timestep Settings

Model	Timesteps	Accuracy (%)	Threshold Type
Baseline	85	86.58	Fixed
Baseline	21	82.80	Adaptive
Pruned	53.08	87.27	Fixed
Pruned	13.04	86.58	Adaptive

Equations

$$u_i^t=\gamma u_i^{t-1}+\sum _j(w_{i,j}{s}_j^{t-1})+b_i-s_i^{t-1}\ast V_{th},$$

$$s_i^t=\hat{\Theta }(u_i^t,V_{th})$$

$$\hat{\Theta }(u_i^t,V_{th})=H(u_i^t-V_{th})$$

$$\frac{\partial {\mathcal{L}}_{pred}(a^{\ast })}{\partial \theta }=-\frac{\partial {\mathcal{L}}_{pred}(a^{\ast })}{\partial a^{\ast }}(J_{g_{\theta }}^{-1}{|}_{a^{\ast }})\frac{\partial f_{\theta }(a^{\ast })}{\partial \theta },$$

$$\arg \underset{h,n}{\min }{\mathcal{L}}_{pred}(h,n)\text{s.t.}M(h,n)\le M_c$$

$$\arg \underset{h,n}{\min }\left(\sum _{i\in S(h)}{\mathcal{I}}_{ii}+\sum _{j\in S(n)}{\mathcal{I}}_{jj}\right)$$

$$I(h_i)={\mathcal{I}}_{ii}\cdot AS{R}_i$$

$$I(n_j)={\mathcal{I}}_{jj}\cdot AS{R}_j$$

$$\arg \underset{h,n}{\min }\left(\sum _{i\in S(h)}I(h_i)+\sum _{j\in S(n)}I(n_j)\right)$$

$$t^l=\lfloor \frac{b^{c^l}}{{\max }_j{b}^{c^j}}\cdot T_{conv}\rfloor$$

$$S_l^i\sim \text{Bernoulli}(a_{l-1}^i))$$

$${\mathcal{L}}_{total}={\mathcal{L}}_{pred}(h,n,t)+\lambda \cdot M(h,n,t)$$

$${\mathcal{L}}_S(a_1^{\ast },\dots ,a_L^{\ast })=\sum _{l=1}^L||a_l^{\ast }|{|}_p$$

$$\text{Normalized \#C}=\frac{{\sum }_{l=2}^{N-1}{a}_l\times {\text{Layer ACs}}^{l+1}}{\sum \text{Layer ACs}}$$

$$M(h,n,t)=\sum _{l=1}^L{\text{ACs}}_h^l+{\text{ACs}}_n^l$$

$${\text{ACs}}_h^l={\text{ACs}}_{\text{QKV}}^l+{\text{ACs}}_{\text{Attn}}^l+{\text{ACs}}_{\text{FC}}^l$$

$${\text{ACs}}_{\text{QKV}}^l={\text{ACs}}_{\text{Q}}\times t_Q^l+{\text{ACs}}_{\text{K}}\times t_K^l+{\text{ACs}}_{\text{V}}\times t_V^l$$

$${\text{ACs}}_{\text{Q}}={\text{ACs}}_{\text{K}}={\text{ACs}}_{\text{V}}=N\times D\times H$$

$${\text{ACs}}_{\text{Attn}}^l=2\times N\times N\times H\times t_{\text{Attn}}^l$$

$${\text{ACs}}_{\text{FC}}^l=N\times D\times H\times t_{\text{FC}}^l$$

$${\text{ACs}}_n^l=N\times D\times t_{\text{Inter}}^l+N\times D\times t_{\text{Output}}^l$$

Full text

I Introduction

Large Language Models (LLMs) have revolutionized the field of natural language processing, providing remarkable advances in tasks such as text generation and machine translation. LLMs, like GPT [1] and BERT [2], have set new benchmarks in terms of accuracy and performance, demonstrating the vast potential of deep learning in understanding and generating human language. However, the large number of parameters and high computational demands translate to considerable energy consumption, limiting their feasibility for real-time and low-power applications. To address these challenges, several compression methods have been developed, including quantization [3, 4], knowledge distillation [5, 6], pruning [7, 8, 9], low-rank factorization [10, 11], and other techniques. These methods aim to reduce the model size of LLMs without significantly compromising their performance. From a complementary compute perspective, Spiking Neural Networks (SNNs) [12], which mimic the spike-based information processing found in biological neural systems, have emerged as a promising alternative for resource-efficient computing [13]. Integrating LLMs with SNNs presents a compelling opportunity to utilize the strengths of both architectures, combining the high accuracy and language understanding capabilities of LLMs with the power efficiency and real-time processing capability of SNNs. Recent works showcase the benefits of using spiking architectures in LLMs such as SpikingBERT [14], Spikformer [15], SpikeLLM [16], among others.

Although spike-based LLMs can reduce power consumption through the inherent sparsity of spiking neurons, the model still requires a substantial amount of operations and retains a large number of parameters. As a result, the overall computational cost still remains high, limiting the scalability and efficiency of spike-based LLMs in practical applications. Further, unlike traditional neural networks, which process information in a single forward pass, SNNs require multiple timesteps to achieve comparable accuracy. Hence there is an unmet need to explore compression techniques catered for spiking LLMs that consider their unique spiking and temporal properties.

Compression techniques (like quantization, pruning, and knowledge distillation), specifically designed for SNNs, have been explored in Refs. [17, 16, 18, 19, 20, 14]. Pruning, in particular, has proven to be an effective strategy for reducing both the model size and inference latency of spiking neurons [21, 22]. In non-spiking models, pruning typically removes redundant or less critical neurons and attention heads, streamlining the network architecture and significantly reducing the number of active parameters. Unlike non-spiking networks, pruning in SNNs must also consider the event-driven and temporal nature of spiking activity, where reducing spiking events and simulation timesteps directly lead to lower energy consumption and latency. Therefore, temporal pruning can be introduced to further optimize SNNs [20, 23]. However, the vast majority of these works remain limited to simple vision based datasets and convolutional architectures. Although initial work [14, 19, 16] has explored quantization and knowledge distillation in spiking LLMs, pruning remains relatively understudied, especially in the context of developing a cohesive framework combining various compression techniques.

The specific contributions of our work are as follows:

(i) Revisiting Spatio-Temporal Pruning in the Spiking LLM Setting: We rethink and extend conventional spatial and temporal pruning strategies by adapting them specifically for SNNs. We propose a principled two-stage pruning framework for Spiking LLMs that incorporates both post-training and retraining phases. In the post-training phase, spatial pruning is guided by a combination of Fisher Information Matrix (FIM), average spiking rate (ASR), and attention output to identify and remove redundant neurons and attention heads. Concurrently, temporal pruning is performed using Principal Component Analysis (PCA) to dynamically allocate timesteps based on the importance of each layer. Furthermore, we integrate PCA analysis for temporal pruning and a computational cost penalty into the loss function for spatial pruning during the retraining phase to enforce energy-aware fine-tuning. This two-stage design ensures that the model achieves high computational efficiency without compromising task performance. Notably, this is the first work to jointly explore the combination of spatial and temporal pruning for large-scale spiking LLMs. This unified pruning strategy is specifically designed to exploit the spatial and temporal characteristics of spiking LLMs.

(ii) Performance Evaluation of Compressed Spiking LLMs on the GLUE Benchmark: We present the first training framework integrating spatial pruning, temporal pruning, extreme quantization and knowledge distillation for Spiking LLMs evaluated on the large-scale GLUE benchmark. Prior SNN compression works have primarily focused on small-scale datasets, limiting their applicability. Our method demonstrates strong scalability and practical feasibility across multiple GLUE tasks, making it a more viable and efficient solution for deploying Spiking LLMs in real-world, energy-constrained environments.

II Related Works

Spiking LLMs: SNNs have emerged as a promising approach to develop compute-efficient LLMs. Spikeformer [15] is the first work to use the spike-based self-attention mechanism, marking a significant step in this direction. SpikeGPT [24] provides the first demonstration of language generation enabled by direct training of SNNs. SpikingBERT [14] uses the average spiking rate of neurons at equilibrium to train the model using implicit differentiation techniques. By avoiding the need for surrogate gradients or explicit time unrolling during model training, this approach enables more efficient and scalable training, providing a solid foundation for further enhancements such as quantization and pruning. We therefore choose SpikingBERT training framework as our baseline for exploring spatio-temporal pruning strategies.

Compression of Spiking LLMs: Compression techniques for Spiking LLMs remain in the nascent stage, with works primarily targeting quantization and knowledge distillation. Refs. [16, 19] investigated quantization methods for Spiking LLMs to reduce weight precision. Knowledge distillation methods, explored in Refs. [25, 14], train smaller Spiking LLMs using larger pre-trained models as teachers. This work primarily focuses on developing a cohesive training framework for Spiking LLMs that incorporates pruning strategies along with extreme quantization and knowledge distillation.

Pruning Methods: Given a trained model, the pruning workflow typically consists of two key stages: applying a pruning method to remove less essential parameters or connections, and fine-tuning the pruned model to restore and improve its performance. Most works focus on the post-training stage [9, 26, 27]. To guide the pruning process more effectively, some approaches leverage the Hessian matrix of the loss function with respect to the model parameters. The weights associated with smaller eigenvalues of the Hessian have less impact on the model’s performance and can be pruned with minimal loss. The technique has been explored and proven effective in prior works [28, 29] for efficiently pruning less critical weights. However, the computational complexity of these methods have limited their scalability to larger models. To address this limitation, Refs. [7, 30] utilized FIM [31] as an approximation of the Hessian to perform post-training pruning. In this work, we adopt this Fisher-based approach and extend it with additional techniques tailored for spiking language models, incorporating spiking-specific metrics, such as average spiking rate, to guide the pruning process.

Pruning Methods for SNNs: Pruning algorithms have been explored in SNN literature. Ref. [32] proposed a gradient-based rewiring pruning method to adjust the network structure by rewiring connections dynamically. In Refs. [22, 33], weight connections are pruned based on the membrane potential. Refs. [34, 35] proposed adaptive weight pruning for SNNs. Ref. [36] presented a method for an unstructured pruning framework with energy consumption constraints. Ref. [20] explored the use of PCA to achieve structured pruning for SNNs, a method also used in ANNs [37]. These works all focus on spatial pruning in SNNs, but unlike ANNs, SNNs have an additional critical dimension: the temporal dimension, which also plays an important role in their efficiency and performance. Ref. [20] proposed a pruning framework focusing on reducing temporal redundancies during the training phase. This approach starts with an SNN trained over multiple timesteps and iteratively reduces the number of timesteps uniformly across the entire network. A similar strategy is also proposed in [38]. However, this uniform approach disregards the varying importance of different layers, which could potentially benefit from a more layer-specific temporal pruning strategy. Ref. [39] investigated the temporal information dynamics in SNNs, revealing that shallow layers exhibit higher importance during the later timesteps, while deeper layers demonstrate significant importance in earlier timesteps. This insight provides the motivation to explore layer-specific temporal pruning. Ref. [40] proposed a dynamic timestep approach that utilizes layer-wise binary masks to reduce inference latency and energy consumption. However, when an early layer is masked, its output becomes zero, which prevents later layers from receiving any input. This can significantly degrade the performance of deeper layers, which are typically more critical for final decision-making. As a result, the effectiveness of this approach is limited, particularly in deep architectures or more complex tasks, and has so far only been demonstrated on small-scale vision benchmarks. This motivates us to develop a more flexible and scalable approach that allows each layer to allocate its own number of timesteps independently, without being constrained by earlier layers. In the domain of spatio-temporal pruning, Ref. [23] integrates both spatial and temporal pruning for Spiking Transformers. However, similar to prior discussions, temporal pruning is applied uniformly across all layers without adaptation to layer-specific dynamics, and the spatial pruning relies solely on spike accumulation values, limiting scalability and flexibility for more complex tasks or architectures. In summary, the vast majority of pruning works explored for SNNs remain limited to small scale tasks and simple convolutional architectures, with only limited works targeting spatio-temporal pruning. Our work overcomes these critical challenges to develop a unified spatio-temporal pruning strategy for Spiking LLMs equipped with quantization and knowledge distillation.

III Methods

In this section, we first introduce the spiking neuron model in SNNs and the associated training technique considered in this work. Subsequently, we present a comprehensive two-stage pruning methodology designed specifically for Spiking LLMs. As shown in Fig. 1, our approach integrates both spatial and temporal pruning techniques to enhance computational efficiency and energy savings, while maintaining model performance. The first stage focuses on post-training pruning. For spatial pruning, we utilize the FIM, ASR, and attention output to identify the redundant neurons and attention heads. This accelerates the model by reducing the number of required operations and enhances energy efficiency by reducing the number of spikes processed during inference. Since the output of each SNN layer is binary spikes instead of continuous activation values, computation involves only accumulation operations (ACs). Therefore, we use the number of ACs to quantify computational cost. For temporal pruning, we apply PCA analysis to dynamically adjust the number of timesteps across layers, reducing temporal redundancy. In the second stage, we refine the model via fine-tuning with pruning. For spatial pruning, we add a penalty term based on ACs to the loss function to reduce overall computational load. For temporal pruning, we continue to use PCA analysis to dynamically reallocate timesteps across layers. Further details on justification for using FIM, ASR, and attention output during spatial pruning and PCA during temporal pruning will be provided in later sections.

For the training process, we adopt the same approach used in SpikingBERT[14] (see Appendix A for details on architecture), which leverages implicit differentiation to compute gradients at a steady-state convergence of the ASR of neurons rather than using Backpropagation Through Time (BPTT). This approach significantly reduces the computational overhead and memory requirements. Additionally, the steady-state ASR convergence of neurons in linear as well as non-linear layers of the SpikingBERT model enables the usage of knowledge distillation, where a pre-trained BERT model serves as the teacher to guide the training of the spiking student model. Quantization of model weights can also be easily integrated into the training framework [19].

III-A SpikingBERT Training Framework

We first introduce the discrete time dynamics of an LIF-based spiking neuron, which is the fundamental component of our SpikingBERT framework,

[TABLE]

[TABLE]

where, $u_{i}^{t}$ is the membrane potential of neuron $i$ at time $t$, $w_{i,j}$ is the weight connecting the pre-synapse neuron $j$ and post-synapse neuron $i$, $s_{j}^{t-1}$ denotes the input spiking signal of pre- synapse neuron $j$ at time $t-1$, $b_{i}$ indicates a bias term, $\gamma$ is the leaky constant and $V_{th}$ is the spiking threshold. $\hat{\Theta}$ is the non-differentiable spiking operation with an underlying Heaviside function, which is defined as,

[TABLE]

where, ${H(x)=\begin{cases}0&\text{if }x<0\\ 1&\text{if }x\geq 0\end{cases}}$. Subtraction serves as the reset operation for $u$ in neurons that spiked.

During training, SpikingBERT leverages implicit differentiation at equilibrium [41], to overcome the non-differentiability of spike-based dynamics. This approach computes gradients based solely on the final converged steady-state of ASR, which is reached after $T_{conv}$ timesteps of convergence. The gradient of the loss $\mathcal{L}_{pred}$ with respect to parameters $\theta$ is computed as:

[TABLE]

where, $a$ denote the ASR value, $a^{*}$ is the steady-state ASR value, $f$ is the steady-state equation of ASR, $g_{\theta}(a)=f_{\theta}(a)-a$, $J^{-1}_{g_{\theta}}$ is the inverse Jacobian of $g_{\theta}$ when $a=a^{*}$, i.e., at equilibrium. Instead of relying on backpropagation through time (BPTT), which requires storing the entire sequence of intermediate states, this approach avoids temporal unrolling by computing gradients at equilibrium, thereby significantly reducing memory consumption and enabling scalable training.

III-B Post-training Pruning

III-B1 Spatial Pruning

To explicitly control the computational cost of the pruned model, we introduce a constraint on the total number of accumulation operations (ACs), which reflects the energy and latency requirements of spiking architectures. The detailed computation of ACs is provided in Appendix B. Since only spatial pruning is considered in this case, all layers operate with the same number of timesteps. Therefore, the timestep variable $t$ can be omitted from the ACs calculation (see Appendix B). In the post-training spatial pruning stage, we determine the optimal binary masks for pruning while satisfying this ACs constraint. We employ two types of masks: neuron masks and head masks. The neuron mask is applied in the feedforward layers to selectively prune less significant neurons, and the head mask is used in the multi-head attention layers to prune redundant attention heads. The pruning problem is formulated as follows:

[TABLE]

where, $\mathcal{L}_{pred}$ is the prediction loss function, $h\in\{0,1\}^{d_{h}}$ is the head mask, and $d_{h}$ is the number of attention heads. Similarly, $n\in\{0,1\}^{d_{n}}$ is the neuron mask, where $d_{n}$ is the number of intermediate neurons. $M$ represents the ACs of the pruned network, and $M_{c}$ is the target ACs constraint. By constraining the total number of ACs, we ensure that the pruned network remains within a desired energy budget. The selection criteria for $M_{c}$ is further explored in our ablation studies (see Section IVD).

However, solving this constrained optimization directly is intractable due to the binary, non-differentiable nature of the mask variables. To address this, we leverage a diagonal approximation of the FIM following Ref. [7]. The FIM provides a second-order approximation of the loss increase when a parameter is removed, quantifying its importance to the model. Thus, instead of minimizing the loss directly, we minimize the cumulative Fisher scores of the pruned units, this leads to the following final objective:

[TABLE]

where, $\mathcal{I}$ is Fisher Information Matrix of gradient of masks, $S(h):=\{i\mid{h}_{i}=0\}$ and $S(n):=\{j\mid{n}_{j}=0\}$. The diagonal elements of the FIM represent the individual sensitivity of each mask with respect to the loss function. In our setting, we treat each of these elements as an importance score: a larger diagonal value indicates that the corresponding unit has a greater influence on the loss, while a smaller value suggests it can be safely pruned. This allows us to rank neurons and attention heads based on their estimated contribution to model performance and selectively retain only the most critical ones.

Unlike conventional neural networks, SNNs communicate using discrete spikes, resulting in inherently sparse and temporally dynamic activity. Motivated by this spiking-specific characteristic, we design a post-training pruning framework that combines traditional sensitivity-based measures, like the FIM, with spike-driven activity statistics, such as ASR and layer-wise timestep usage. The ASR is a direct measure of how frequently each neuron contributes to the computation, making it an essential factor in determining neurons’ importance. Therefore, the final importance score for each neuron mask should be a combination of both the FIM and the corresponding ASR. Similarly, for the attention mechanism, a head with larger output values should have a more substantial impact on the outputs of the fully connected layer, as well as the entire transformer block. We use the ASR of the attention outputs to assess the importance of each head, reflecting how frequently a head contributes to computation. This idea conceptually aligns with prior work in non-spiking models, where head importance is estimated by accumulating the absolute values of attention outputs, as shown in Ref. [9]. The final importance of neuron masks and head masks can be summarized as follows:

[TABLE]

[TABLE]

where, $\mathcal{I}_{ii}$ and $\mathcal{I}_{jj}$ are the Fisher information importance scores for head $i$ and neuron $j$ respectively. $ASR_{i}$ and $ASR_{j}$ denote the average spiking rate of the attention output of head $i$ and average spiking rate of neuron $j$ respectively. By combining FIM and ASR, these scores provide a more comprehensive measure of unit importance in the spiking domain. Consequently, the pruning objective in Eqn. 6 becomes:

[TABLE]

After computing the per-unit importance scores, neurons and attention heads are ranked accordingly. The lowest-ranked units are then pruned to satisfy the ACs constraint. Following this initial pruning, we perform a mask search and refinement phase, where a greedy algorithm iteratively swaps pruned and unpruned units within each layer. This refinement step helps account for intra-layer dependencies and enhances the effectiveness of the pruning strategy. This approach is inspired by Ref. [7].

III-B2 Temporal Pruning

In the post-training stage, we also introduce a temporal pruning strategy to further improve efficiency in Spiking LLMs. In traditional SNNs, every timestep contributes to the forward pass, which can lead to inefficiencies, especially when many timesteps carry redundant or low-importance information. Moreover, not all layers may require the same temporal resolution. We hypothesize that more important layers in a network require more time unrolling. To address this, we propose a temporal pruning strategy that dynamically adjusts the number of timesteps per layer based on its temporal complexity. To quantify this temporal complexity, we apply Principal Component Analysis to the layer-wise activations across timesteps. This approach is inspired by Ref. [42], which employed PCA to infer the intrinsic dimensionality and redundancy of neural network layers in the context of adaptive quantization. Other works have used PCA-based structural pruning technique in Refs. [37, 20]. Distinct from earlier approaches, our method targets the analysis of temporal dynamics rather than spatial redundancy. In our work, we record the ASR of each layer at every timestep, then apply PCA to this temporal sequence to determine the number of principal components required to explain the majority of the temporal variance. This effectively captures the complexity of its temporal dynamics. Subsequently, we map the PCA-derived importance of each layer to the number of timesteps allocated per layer using a nonlinear function to emphasize differences in temporal complexity across layers. The mapping function can be formulated as:

[TABLE]

where, $t^{l}$ denotes the number of timesteps allocated to layer $l$, $c^{l}$ is the PCA-derived importance score of layer $l$, $b$ is a scaling factor (base) for the power function, and $T_{conv}$ is the number of timesteps required for convergence. To ensure that layers with higher importance scores run for more timesteps, the scaling factor $b$ must be greater than 1. This mapping ensures that layers with higher temporal complexity are assigned more timesteps, preserving their dynamic representational capacity. Conversely, layers with lower temporal complexity receive fewer timesteps, allowing for more aggressive temporal pruning without significantly compromising performance.

It is important to discuss the architectural dataflow in the context of our proposed temporal pruning. From a high-level hardware implementation perspective, there are two primary approaches to processing Spiking LLMs: (1) complete spatial unrolling, where all layers are simultaneously mapped onto the chip, and (2) resource-constrained sequential mapping, where layers are processed one at a time due to hardware limitations. The first approach allows parallel computation but requires significant hardware resources, which is often impractical for large-scale models. The second one is a more realistic approach that reuses limited compute resources by sequentially allocating them to different layers, as shown in Fig. 2. The left panel depicts the first structure of Spiking LLMs, where spiking activity propagates sequentially through all layers, with each layer processing the same number of timesteps during the forward pass. The right panel shows the second approach, which facilitates the reuse of limited hardware resources by sequentially allocating them to different layers. This is relevant for our scenario to implement dynamic timestep allocation to various layers of the network. However, it necessitates storing the complete set of output spikes from all timesteps of a given layer for use by subsequent layers, which leads to significant memory overhead. To solve this problem, we propose stochastic spiking LLMs, which adopts a stochastic approach to generate the input spike trains to each layer. Instead of storing all the spikes at each timestep, which would be memory-intensive, we utilize the average spiking rate of the previous layer to stochastically generate spikes for the current layer. The equation for generating the input spikes in layer $l$ based on the ASR of the previous layer $a_{l-1}$ can be formulated as:

[TABLE]

where, $S_{l}^{i}\in\{0,1\}^{t^{l}}$ is the binary spike train for neuron $i$ in layer $l$, with the length $t^{l}$ being the number of timesteps allocated for layer $l$. Each element of the spike train is independently sampled from a Bernoulli distribution with the probability $a_{l-1}^{i}$, the ASR of neuron $i$ in the previous layer $l-1$. By sampling from this Bernoulli distribution, we effectively simulate the spiking behavior without the need to retain detailed spike histories, thus saving memory and computational power. This is possible since we rely on rate encoding for our SNN operation.

III-C Pruning During Retraining

Spatial pruning and temporal pruning can also be incorporated effectively in the retraining stage to further enhance the compression efficiency. We leverage the head and neuron masks, as well as the layer-wise timestep allocations, obtained from the post-training pruning stage. In order to preserve spatial pruning while adhering to computational cost constraints, we incorporate a computational cost penalty into the loss function. Similar to the approach outlined in Ref. [36], we modify the loss function to include a penalty term proportional to the number of ACs required by the pruned network. Since the number of ACs in SNNs depends on both spatial and temporal factors, the ACs should also consider the number of timesteps allocated to each layer during temporal pruning. The new loss function is given by the following equation:

[TABLE]

where, $\mathcal{L}_{pred}(h,n,t)$ is the standard prediction loss of the pruned model with head mask $h$ and neuron mask $n$. $M(h,n,t)$ represents the total number of ACs required by the pruned network, which is a function of the head mask $h$, neuron mask $n$, and timesteps list $t=[t^{1},t^{2},\ldots,t^{l}]$, where each element $t^{l}$ denotes the number of timesteps assigned to layer $l$. $\lambda$ is a hyperparameter controlling the trade-off between minimizing the loss and reducing the computational cost, effectively balancing performance and computational efficiency. Due to the binary nature of the mask, there is a non-differentiability problem [36]. To enable gradient-based training, we adopt a surrogate gradient where we approximate the binary mask using a sigmoid-based function.

We employ a staged optimization strategy: During the initial training epochs, we incorporate the ACs penalty in the loss function to encourage low computational usage. Once the model stabilizes, we switch back to the original loss function to recover and preserve accuracy. This staged optimization helps balance efficiency with performance. Temporal pruning also utilizes post-training results as the initialization point. At regular intervals during fine-tuning, we perform PCA analysis on updated spiking activity to reallocate the number of timesteps for each layer, allowing the model to adapt to its evolving temporal dynamics.

III-D Additional Optimizations

We augment our training framework with a few additional optimizations to improve the sparsity and energy efficiency of our Spiking LLMs.

The threshold potential is pivotal in spike generation and serves as a crucial hyperparameter in SNN design. Unlike conventional approaches where a fixed threshold is applied uniformly across all layers, we treat the threshold $V_{th}$ in Eqn. 1 as a learnable parameter that is independently fine-tuned for each layer during training. Therefore, the network can allow each layer to adapt its spiking behavior. This flexibility is especially important in deep spiking networks, where different layers may require different levels of sparsity and responsiveness.

Additionally, sparsity is another key characteristic of SNNs, contributing directly to their energy efficiency. While our pruning strategies already reduce computational cost, promoting sparsity offers an additional aspect of optimization. To explicitly encourage sparse spiking activity during training, we introduce a novel loss function, given as $\mathcal{L}_{total}=\mathcal{L}_{pred}(h,n)+\lambda\cdot M(h,n,t)+\eta\cdot\mathcal{L}_{S}$. The hyperparameters $\lambda$ and $\eta$ control the trade-off between task performance, computational cost, and spiking sparsity. The activity loss function $\mathcal{L}_{S}$ is given as,

[TABLE]

where, $||a^{*}_{l}||_{p}$ is the $p$-norm of the converged ASR for the output of the $l$-th spiking encoder layer. In our case, we use $p=2$, corresponding to the Euclidean norm. By incorporating $\mathcal{L}_{S}$, the model achieves improved sparsity without compromising pruning effectiveness or task performance.

IV Experiments

IV-A Datasets

We evaluated our model’s performance using a diverse set of text classification tasks from the General Language Understanding Evaluation (GLUE) benchmark [43]. Specifically, we employed the Quora Question Pair (QQP) and Microsoft Research Paraphrase Corpus (MRPC) datasets to assess the model’s capability in handling similarity and paraphrase detection tasks. For inference-oriented evaluations, we utilized the Multi-Genre Natural Language Inference (MNLI) and Question-answering Natural Language Inference (QNLI) datasets to measure the model’s ability to perform reasoning across various contexts and question-answering scenarios. For sentiment analysis, we leveraged the Stanford Sentiment Treebank (SST-2) dataset to evaluate the model’s performance on single-sentence sentiment classification tasks. For all experiments, we set the maximum sequence length to 128 tokens.

IV-B Experimental Setup

The experiments were conducted on a system equipped with 8 NVIDIA RTX A5000 GPUs, each with 24 GB of memory. For each run, we employed PyTorch DataParallel using two GPUs. In our experiments, we use task-specific 1-bit SpikingBERT model with 4 encoder layers as the baseline for each individual task. All models are initialized from a full-precision SpikingBERT [14]. We quantize the model weights to 1-bit following the approach in Ref. [19]. To mitigate any performance degradation caused by quantization, we apply intermediate-layer knowledge distillation using the full-precision teacher model, as described in [19]. Finally, we fine-tune the quantized model on the corresponding task-specific dataset to obtain the final model. The detailed hyperparameter settings are listed in Table I.

To quantify the event-driven operation of SNNs, we incorporated a Normalized Operations metric. In contrast to ANNs, where all neurons are active during each forward pass, SNNs only activate fan-out neurons when a spike occurs. Therefore, the number of effective operations in SNNs depend not just on the model architecture, but also on the spiking activity during inference, typically characterized by the average spiking rate. We adopt a normalized operations metric as used in [44, 14]. The normalized computational factor is defined as:

[TABLE]

where, $N$ is the total number of layers (where each sublayer within the encoder blocks is considered separately), $a_{l}$ is the average spiking rate of layer $l$, calculated as the total number of spikes generated by all neurons in the layer divided by the number of neurons and the number of timesteps. $\text{Layer ACs}^{l+1}$ represents the number of AC operations in layer ${l+1}$. In our experiments, we employ the Normalized #C metric to assess the efficiency improvements achieved by our pruning strategy and compare the results against the baseline model as well as other optimization approaches.

IV-C Results

The accuracy (Acc) and ACs of the 1-bit quantized SpikingBERT model using various pruning methods are summarized in Table II. In our post-training spatial pruning approach, we evaluate different importance scores to determine which neurons and heads to prune. Post-Training (FIM), which uses only the Fisher Information Matrix as the importance score, resulted in the lowest accuracy among the methods tested. To improve accuracy, Post-Training (FIM + ASR of Neurons) and Post-Training (FIM + ASR of Attn) consider the ASR of neurons and the attention output of each head in the importance scores, separately. The results show that they both can improve accuracy. Post-Training (FIM + ASR) combines these three scores, which achieves the highest accuracy. After obtaining the pruning masks of neurons and attention heads from the post-training spatial pruning stage, we apply a retraining pruning stage to further improve the model’s performance. During retraining, we only train for a small number of epochs. Despite the short retraining period, this approach results in a notable improvement in accuracy.

For temporal pruning, in the retraining stage, the model is fine-tuned with the newly constrained temporal structure. By allowing the network to adapt to the reduced timestep configuration, retraining helps the model to re-optimize its parameters and adjust its temporal processing. Two-stage pruning integrates spatial and temporal pruning. We use the retrained spatially pruned models to perform temporal pruning (post-training and retraining). Additionally, we also incorporate the sparsity improvement optimization strategy by including the activity loss in the loss function during this final step.

Further, by incorporating the adaptive thresholding mechanism with our pruned model during the final retraining stage, the number of required timesteps can be further reduced while maintaining model accuracy. In our experiments, we used the pruned timestep list and scaled it down to a lower range. We then retrained the model using this new timestep list in combination with the adaptive threshold mechanism. The adaptive threshold allows layers to dynamically adjust their firing behavior in response to this scaling. The accuracy, latency, and computational cost of various 1-bit quantized models are presented in Table  III. It is worth mentioning here that the latency is reported as the average number of timesteps across all layers. Compared with the baseline model, our pruning method significantly reduces both computational cost and latency. Specifically, the total Normalized #C metric drops to approximately 58.51% to 84.19% of the baseline values, while the mean latency is reduced to around 28.86% to 50.91%. The adaptive thresholding mechanism can further reduce the latency and computational cost, down to only 89.85% to 95.22% of the baseline #C and 79.86% to 84.66% of the baseline latency values. This demonstrates the effectiveness of our approach in optimizing Spiking LLMs for energy-efficient and real-time inference scenarios.

IV-D Ablation Studies

Effect of Activity Loss: We performed experiments to evaluate the impact of the activity loss term $\mathcal{L}_{S}$ on model behavior. The steady-state behavior of the mean layer-specific ASR in the optimized model and the layer-wise reduction in neuronal activity due to activity loss $\mathcal{L}_{S}$, is shown in Fig. 3 and 3. We see that $\mathcal{L}_{S}$ leads to a notably lower ASR. This reduction indicates that the activity loss effectively suppresses redundant neuronal activity, leading to a more energy-efficient network.

ACs Constraint Selection: To evaluate the impact of computational constraints in post-training spatial pruning, we conducted experiments to explore various ACs budget settings. Specifically, we varied the ACs constraint from 0.4 to 1.0, which means that the pruned model is restricted to use only 40% to 100% of the original model’s accumulation operations, as shown in Fig. 4. Based on the observed trade-off between model performance and computational efficiency, we selected the ACs constraint at 60%. This setting achieves a desirable balance, significantly reducing the computational cost while maintaining acceptable performance.

PCA Scaling Factor Selection: To determine the optimal scaling factor involved in mapping the PCA metric to the layerwise timestep allocation using a power function based non-linearity during temporal pruning, we performed experiments by varying the scaling factor from 1 to 1.2, and evaluating its impact on both the number of computational operations (ACs) and the accuracy (Acc) of the model. As shown in Fig. 4, increasing the scaling factor progressively enforces stricter pruning, thereby reducing the total ACs. However, overly aggressive pruning leads to a significant drop in accuracy. Among the tested values, a scaling factor of 1.02 strikes an effective balance. It can preserve accuracy comparable to the baseline while achieving a notable reduction in ACs.

Comparison of Two-Stage vs. Retraining-Only Spatial Pruning: We evaluated the effectiveness of our two-stage spatial pruning strategy by comparing it with the retraining-only pruning, in which pruning masks are initialized randomly rather than being informed by a post-training phase. However, this approach did not achieve accuracies comparable to those obtained using the two-stage pruning. Furthermore, this approach required more epochs to train. As shown in Fig. 5, the post-training followed by retraining approach achieved faster convergence and attains higher training accuracy compared to retraining from a random mask. The improved efficiency of the two-stage method is attributed to the initial pruning step, which removes less important neurons and attention heads, allowing subsequent retraining to more effectively fine-tune the model within the reduced architecture.

Two-Stage Spatial and Temporal Pruning Strategy: We also designed experiments to investigate the optimal strategy for pruning our model. During post-training pruning, spatial and temporal pruning are applied independently to the baseline model. Once the best spatial masks and temporal timestep lists are identified, we proceed with retraining. Table IV demonstrates that performing a sequential strategy that apply spatial pruning followed by retraining, and then applying temporal pruning with an additional retraining stage, leads to better accuracy compared to jointly retraining with both spatial and temporal pruning applied at once. Without intermediate recovery, the model struggles to cope with the simultaneous reduction in spatial capacity and temporal resolution, making convergence more difficult and increasing the risk of getting trapped in suboptimal local minima. Although our approach involves two retraining stages, the use of post-training pruning significantly reduces the number of epochs required for each retraining phase, making the overall process efficient.

Effect of Adaptive Threshold: To evaluate the impact of adaptive thresholding, we conducted experiments using the baseline model with the adaptive threshold mechanism. Compared to the baseline model without pruning, the adaptive threshold mechanism proves to be more effective when applied to a pruned model with reduced timesteps. As shown in Table. V, we tested different thresholding strategies and timestep configurations using the 1-bit quantized model on the SST-2 dataset. It significantly improves performance in the pruned model under aggressive timestep reduction. This highlights the effectiveness of adaptive thresholding as a compensatory mechanism for pruning and reducing temporal resolution. Another contributing factor is our resource-constrained sequential mapping architecture, which introduces stochasticity during spike generation. This stochastic behavior avoids the typical spike propagation delay associated with sequential information propagation in deep spiking architectures, allowing it to maintain performance even when the number of timesteps is significantly reduced.

V Conclusions

This is the first work to jointly explore scalable spatial pruning, temporal pruning, extreme quantization, and knowledge distillation for Spiking LLMs evaluated on a large-scale NLP benchmark. By leveraging a two-stage pruning strategy, our method achieves substantial gains in latency and computational efficiency without compromising performance. These results demonstrate the effectiveness of combining SNNs with structured compression techniques to build highly efficient and deployable language models for energy-constrained applications. Future work can include extending the application of our pruning framework to other architectures and tasks, as well as exploring additional compression methods specifically tailored for spiking LLMs. By advancing these techniques, we aim to further bridge the gap between high-performance language models and energy-efficient computation on edge devices.

Acknowledgments

This material is based upon work supported in part by the U.S. National Science Foundation under award No. CAREER #2337646, CCSS #2333881, CCF #1955815, and EFRI BRAID #2318101.

Appendix A SpikingBERT Architecture

The architecture of the SpikingBERT model used in our work is shown in Fig. 6. This design retains the core structure of the Transformer encoder, with modifications to enable spike-based computation. In this architecture, each encoder block comprises of six spiking sublayers: (1) Key projection layer, (2) Value projection layer, (3) Spiking Attention layer, (4) the first Add & Norm linear layer (self-attention output layer), (5) the Feed-Forward layer (intermediate layer), and (6) the second Add & Norm linear layer (encoder output layer). The head mask is applied to sublayer (3), while the neuron mask is applied to sublayers (4) and (5). In our implementation, the Key and Value components of the attention mechanism are converted into spike-based representations, while the Query remains in its real-valued form (based on prior formulations [14]). We therefore consider the Query projection and subsequent attention computation to be integrated within the spiking attention module. The attention output is subsequently converted into spikes for downstream processing. Additionally, we incorporated LIF neurons to generate spikes, along with residual connections and layer normalization modules, similar to the original BERT architecture.

Appendix B Computation of Number of Accumulation Operations

We now detail the computation of the number of accumulation operations (ACs), denoted as $M(h,n,t)$, as introduced in Eqn. 5 and Eqn. 12. Since each layer operates for a different number of timesteps, the equation can be expressed as:

[TABLE]

where, $L$ is the total number of encoder layers, $\text{ACs}_{h}^{l}$ represents the ACs for all attention heads in layer $l$, and $\text{ACs}_{n}^{l}$ represents the ACs for all neurons in layer $l$. $\text{ACs}_{h}^{l}$ is composed of three components: the ACs for the query, key, and value operations ($\text{ACs}_{\text{QKV}}^{l}$), the ACs for the attention computation ($\text{ACs}_{\text{Attn}}^{l}$), and the ACs for the fully connected self-attention output layer ($\text{ACs}_{\text{FC}}^{l}$). This relationship is summarized by the following equation:

[TABLE]

The ACs for the sum of the linear transformations of the query, key, and value operations can be represented as:

[TABLE]

where, $\text{ACs}_{\text{Q}}$, $\text{ACs}_{\text{K}}$, $\text{ACs}_{\text{V}}$ are the ACs for the query, key and value respectively, scaled by their corresponding timesteps $t_{Q}^{l}$, $t_{K}^{l}$, and $t_{V}^{l}$. The ACs for the query, key, and value operations are equivalent and can be computed as:

[TABLE]

where, $N$ is the sequence length, $D$ is the hidden size, and $H$ is the dimension of each attention head which is equal to $D/{d^{h}}$. The ACs for the attention computation and for the fully connected self-attention output layer are given by the following equations:

[TABLE]

[TABLE]

where, $t_{\text{Attn}}^{l}$ and $t_{\text{FC}}^{l}$ denote the timesteps of the self-attention layer and self-attention output layer, respectively. In our architecture, since the Query projection is implemented within the spiking attention layer, we set $t_{Q}^{l}=t_{\text{Attn}}^{l}$. The attention layer contains two dot product operations for each head in layer $l$, first computing attention scores by comparing queries with keys, and the second for generating the weighted output by applying the attention scores to the values.

The accumulation operations for all neurons in the $l^{th}$ layer are given by:

[TABLE]

where $t_{\text{Inter}}^{l}$ and $t_{\text{Output}}^{l}$ denote the number of timesteps allocated to the intermediate and output layers, respectively.

Post-training spatial pruning requires a fixed architectural constraint to guide the pruning of attention heads and neurons. Since the pruning masks are determined after training without further tuning of spiking behavior, incorporating ASR introduces variability and undermines the determinism of the pruning process. Therefore, to maintain consistency and control over the computational budget during post-training pruning, we rely solely on structural metrics (e.g., number of heads or neurons) rather than dynamic activity-based measures like ASR.