Winner-Take-All Spiking Transformer for Language Modeling
Authors: Chenlin Zhou, Sihang Guo, Jiaqi Wang, Dongyang Ma, Kaiwei Che, Baiyu Chen, Qingyan Meng, Zhengyu Ma, +1 more
Organizations: School of Electronic and Computer Engineering, Shenzhen Graduate School, Peking University · Pengcheng Laboratory · Harbin Institute of Technology · School of Computer Science, Peking University
Spiking Transformers, which combine the scalability of Transformers with the sparse, energy-efficient property of Spiking Neural Networks (SNNs), have achieved impressive results in neuromorphic and vision tasks and attracted increasing attention. However, existing directly trained spiking transformers primarily focus on vision tasks. For language modeling with spiking transformer, convergence relies heavily on softmax-based spiking self-attention, which incurs high energy costs and poses challenges for neuromorphic deployment. To address this issue, we introduce Winner-Take-All (WTA) mechanisms into spiking transformers and propose two novel softmax-free, spike-driven self-attention modules: WTA Spiking Self-Attention (WSSA) and Causal WTA Spiking Self-Attention (CWSSA). Based on them, we design WTA-based Encoder-only Spiking Transformer (WE-Spikingformer) for masked language modeling and WTA-based Decoder-only Spiking Transformer (WD-Spikingformer) for causal language modeling, systematically exploring softmax-free, spiking-driven Transformer architectures trained end-to-end for natural language processing tasks. Extensive experiments on 16 datasets spanning natural language understanding, question-answering tasks, and commonsense reasoning tasks validate the effectiveness of our approach and highlight the promise of spiking transformers for general language modeling and energy-efficient artificial intelligence.
Figures & tables
Figure 1: Overview of WTA Spiking Self-Attention (WSSA) and Causal WTA Spiking Self-Attention (CWSSA). The left shows the softmax-based spiking self-attention in SpikeLM ( Xing et al., 2024b ) . The right shows our WSSA in WE-Spikingformer for masked language modeling and CWSSA in WD-Spikingformer for causal language modeling. The symbol “@” denotes matrix multiplication.
Figure 2: The overview of WE-Spikingformer and WD-Spikingformer. The left shows WE-Spikingformer (WTA-based Encoder-only Spiking Transformer) for spike-based masked language modeling. The right shows WD-Spikingformer (WTA-based Decoder-only Spiking Transformer) for spike-based causal language modeling.
Model
SD
T
MNLI_mm
QQP
QNLI
SST-2
CoLA
STS-B
MRPC
RTE
Avg. (%)
BERT base ( Devlin et al., 2019 )
✗
–
83.4
71.2
90.5
93.5
52.1
85.8
88.9
66.4
79.6
Q2BERT ( Zhang et al., 2020 )
✗
–
47.3
67.0
61.3
80.6
0.0
4.7
81.2
52.7
49.1
BiBERT ( Qin et al., 2022 )
✗
–
67.5
84.8
72.6
88.7
25.4
33.6
72.5
57.4
63.2
SpikeLM ( Xing et al., 2024b )
✗
4
77.2
83.9
85.3
87.0
38.8
84.9
85.7
69.0
76.5
LIF-BERT ( Xing et al., 2024b )
✓
4
35.2
0
50.5
50.9
0
0
81.2
52.7
34.6
PSN-BERT ( Xing et al., 2024b )
✓
4
35.2
0
50.5
50.9
0
6.8
81.2
52.7
34.7
Table 1: The results on the Natural Language Understanding task (GLUE datasets). ”Avg.” denotes ”Average Accuracy ”. ”SD” denotes ”Spike-Driven”. The results of LIF-BERT, PSN-BERT, and SpikeLM are reported in Xing et al. (2024b) . WE-Spikingformer uses NI-LIF neuron, while WE-Spikingformer † adopts the T-LIF neuron. Our model achieves state-of-the-art performance on the GLUE datasets among softmax-free spike-driven transformer models.
Model
SD
E (mJ)
T
ARC-e
ARC-c
BoolQ
HeadQA
OBQA
Avg. (%)
SpikeLLM-7B ( Xing et al., 2024a )
✗
-
4
31.3
23.6
53.8
-
-
-
Qwen-1.5B ( Yang et al., 2024 )
✗
3398.3
-
26.2
26.9
60.3
25.7
27.1
33.2
WD-Spikingformer-0.4B
✓
238.4
4
30.0
22.3
37.8
26.0
26.1
28.4
Table 2: The results on Question Answering Tasks (QAT). ”Avg.” denotes ”Average Accuracy”. ”E (mJ)” means ”Energy consumption (mJ)”. ”T” means time step. ”SD” denotes ”Spike-Driven”.
Model
SD
E (mJ)
T
HellaSwag
PIQA
Winograd
Avg. (%)
SpikeLLM-7B ( Xing et al., 2024a )
✗
-
4
33.9
53.4
51.5
46.3
Qwen-1.5B ( Yang et al., 2024 )
✗
3398.3
-
26.5
53.7
52.8
44.3
WD-Spikingformer-0.4B
✓
238.4
4
25.9
53.4
50.2
43.2
Table 3: The results on Commonsense Reasoning Tasks (CRT). ”Avg.” denotes ”Average Accuracy”. ”E (mJ)” means ”Energy consumption (mJ)”. ”T” means time step. ”SD” denotes ”Spike-Driven”.
Datasets
T
Method
Avg.
GLUE
4
QKTV∗s
55.4
4
Softmax(QKT∗s)V
66.8
4
HardWTA(QKT∗s)V
66.3
QAT
4
M(QKT∗s)V
✗
4
Softmax(M(QKT∗s))V
28.7
4
HardWTA(M(QKT∗s))V
28.4
Table 4: Ablation study for HardWTA on GLUE, QAT, and CRT. M(.) means causal mask for decoder-only spiking transformers. That is, M(QKT∗s)V is the QKTV∗s in causal mask version. Note that ”✗” means the model does not converge in pretraining. ”T” means time step. ”Avg.” denotes ”Average Accuracy ( % )”.
Datasets
T
Method
Avg.
GLUE
4
HardWTA(QKT∗s)V
65.7
4
Top-KWTA(QKT∗s)V
66.3
4
Sparsemax(QKT∗s)V
66.6
8
HardWTA(QKT∗s)V
67.4
4
HardWTA(QKT∗s)V
65.7
2
HardWTA(QKT∗s)V
62.3
Table 5: Ablation study on different WTAs and time steps on GLUE datasets. ”T” means time step. ”Avg.” denotes ”Average Accuracy ( % )”.
Figure 3: Increasing model Parameters for WE-Spikingformer pretraining on (a) Question-Answering Tasks (QAT), and (b) Commonsense Reasoning Tasks (CRT).
Figure 4: Increasing tokens for WE-Spikingformer-1.0B pretraining on (a) Question-Answering Tasks (QAT), and (b) Commonsense Reasoning Tasks (CRT).
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
Model
SD
E (mJ)
T
ARC-e
ARC-c
BoolQ
HeadQA
OBQA
Avg. (%)
SpikeLLM-7B ( Xing et al., 2024a )
✗
-
4
31.3
23.6
53.8
-
-
-
Qwen-1.5B ( Yang et al., 2024 )
✗
3398.3
-
26.2
26.9
60.3
25.7
27.1
33.2
WD-Spikingformer-0.4B
✓
238.4
4
30.0
22.3
37.8
26.0
26.1
28.4
WD-Spikingformer-1.0B
✓
574.7
4
33.8
20.3
37.8
26.4
26.6
29.0
WD-Spikingformer-1.0B ∗
✓
574.7
4
41.9
20.4
42.1
26.7
28.2
31.9
Appendix
Table 6: The results on Question Answering Tasks (QAT). ”Avg.” denotes ”Average Accuracy”. ”E (mJ)” means ”Energy consumption (mJ)”. ”T” means time step. ”SD” denotes ”Spike-Driven”. WD-Spikingformer-1.0B ∗ means WD-Spikingformer-1.0B with 0.5B tokens during pretraining.
Model
SD
E (mJ)
T
HellaSwag
PIQA
Winograd
Avg. (%)
SpikeLLM-7B ( Xing et al., 2024a )
✗
-
4
33.9
53.4
51.5
46.3
Qwen-1.5B ( Yang et al., 2024 )
✗
3398.3
-
26.5
53.7
52.8
44.3
WD-Spikingformer-0.4B
✓
238.4
4
25.9
53.4
50.2
43.2
WD-Spikingformer-1.0B
✓
574.7
4
26.0
55.8
49.1
43.6
WD-Spikingformer-1.0B ∗
✓
574.7
4
26.4
55.7
50.3
44.1
Appendix
Table 7: The results on Commonsense Reasoning Tasks. ”Avg.” denotes ”Average Accuracy”. ”E (mJ)” means ”Energy consumption (mJ)”. ”T” means time step. ”SD” denotes ”Spike-Driven”. WD-Spikingformer-1.0B ∗ means WD-Spikingformer-1.0B with 0.5B tokens during pretraining.
Backbone
Neuron
T
Average accuracy ( % )
Training time ( hours)
WE-Spikingformer
NI-LIF
4
65.7
26
WE-Spikingformer
T-LIF
4
66.3
121
Appendix
Table 8: Spiking neuron comparison for WE-Spikingformer pretraining.
Surrogate Gradient
T
MNLI_mm
QQP
QNLI
SST-2
CoLA
STS-B
MRPC_F1
RTE
Avg. (%)
Softmax
4
73.4
83.8
78.3
85.6
23.7
56.0
76.4
48.4
65.7
ReLU
4
31.8
63.2
50.5
50.9
0
4.3
75.0
47.3
40.4
Appendix
Table 9: Ablation study for Surrogate Gradient on GLUE datasets.
Method
T
MNLI_mm
QQP
QNLI
SST-2
CoLA
STS-B
MRPC_F1
RTE
Avg. (%)
Softmax(QKT∗s)V
4
72.5
84.7
76.0
87.2
24.4
54.5
79.7
55.6
66.8
Top-K WTA (K=1)
4
73.4
83.8
78.3
85.6
23.7
56.0
76.4
48.4
65.7
Top-K WTA (K=3)
4
73.4
84.0
78.8
86.1
24.1
55.9
78.2
49.8
66.3
Top-K WTA (K=5)
4
73.1
84.4
79.4
86.0
24.8
55.4
79.1
50.1
66.5
Appendix
Table 10: Ablation study for Top-K Winner-Take-All on GLUE datasets. Top-K WTA (K=1) is equal to Hard WTA.
Datasets
Time step
Method
Average Accuracy ( % )
CRT
4
HardWTA(M(QKT∗s))V
43.2
4
Top-KWTA(M(QKT∗s))V
42.9
4
Sparsemax(M(QKT∗s))V
43.2
8
HardWTA(M(QKT∗s))V
43.7
4
HardWTA(M(QKT∗s))V
43.2
2
HardWTA(M(QKT∗s))V
40.1
Appendix
Table 11: Ablation study for different WTAs and time steps on CRT.
Spiking neural networks promise brain-inspired neural dynamics, yet native large-scale spiking language-model pretraining remains largely unexplored relative to conversion, distillation, and backbone-adaptation approaches. We diagnose a central obstacle: standard spiking neurons apply the same firing sensitivity to every token, creating dense- or silent-firing regimes that are difficult to resolve with globally tuned parameters alone. This problem is particularly acute in language modeling, where a single sequence contains tokens of vastly different informational roles, making a uniform firing regime poorly matched to the data. We address this with NeuronSpark, a self-attention-free language model built around input-dependent spiking membrane dynamics, where each neuron's decay rate, gain, and firing threshold adapt to the current token content. Together with gated associative memory and adaptive temporal aggregation, this selectivity enables billion-scale pretraining from random initialization without backbone transfer, conversion, or distillation. We relate the neuron to fixed-parameter spiking models and the memory to linear attention. At 1.1B parameters and 2.1B tokens, NeuronSpark reaches a 40.0 eight-task zero-shot average, exceeds comparable-budget Pythia/Mamba checkpoints, retains 95.2%/93.2% of their fully trained scores, and spontaneously develops a fast--slow hierarchy. At 2,048 tokens, parameter-matched VHDL points reduce operation-energy proxy by 90.1%/67.0% and cycles by 94.1%/80.3% versus Transformer/Mamba; a physical FPGA path reaches 24.000s TTFT and 0.082440 TPS.
Spiking Neural Networks (SNNs) offer promising energy-efficient alternatives to large language models (LLMs) due to their event-driven nature and ultra-low power consumption. However, to preserve capacity, most existing spiking LLMs still incur intensive floating-point matrix multiplication (MatMul) and nonlinearities, or training difficulties arising from the complex spatiotemporal dynamics. To address these challenges, we propose BiSpikCLM, the first fully binary spiking MatMul-free causal language model. BiSpikCLM introduces Softmax-Free Spiking Attention (SFSA), eliminating softmax and floating-point operations in autoregressive language modeling. For efficient training, we introduce Spike-Aware Alignment Distillation (SpAD), which aligns ANN teacher and SNN student across embeddings, attention maps, intermediate features, and output logits. SpAD framework allows BiSpikCLM to reach comparable performance to ANN counterparts using substantially fewer training tokens (e.g., only 5.6% of the tokens for the 1.3B model). As a result, BiSpikCLM achieves competitive performance at only 4.16% - 5.87% of the computational cost on natural language generation tasks. Our results highlight the feasibility and effectiveness of fully binary spike-driven LLMs and establish the distillation as a promising pathway for brain-inspired spiking NLP.
Sihang Guo, Chenlin Zhou, Jiaqi Wang +3
Harbin Institute of Technology, Shenzhen, China · School of Electronic and Computer Engineering, Shenzhen Graduate School, Peking University, Shenzhen, China · Peng Cheng Laboratory, Shenzhen, China
The Transformer architecture is widely regarded as the most powerful tool for natural language processing, but due to a high number of complex operations, it inherently faces the issue of high energy consumption. To address this issue, we consider Spiking Neural Networks (SNNs), which are an energy-efficient alternative to conventional Artificial Neural Networks (ANNs) due to their naturally event-driven approach to processing information. However, this inherently makes them difficult to train. Often, many SNN-based models circumvent this issue by converting pre-trained ANNs. More recently, attempts have been made to design directly trainable SNN-based adaptations of the Transformer model structure. Although the results showed great promise, the application field was computer vision. Moreover, the proposed model incorporates only encoder blocks. In this paper, we propose SpikeDecoder, a fully SNN-based implementation of the Transformer decoder block, for applications in natural language processing. In a series of experiments, we analyze the impact of exchanging different blocks of the ANN model with spike-based alternatives to identify trade-offs and significant sources of performance loss. We further investigate the role of residual connections and the selection of SNN-compatible normalization techniques. Besides the work on the model architecture, we formulate and compare different embedding methods to project text data into spikes. Finally, we demonstrate that our proposed SNN-based decoder block reduces the theoretical energy consumption by 87% to 93% compared to the ANN baseline.
Claas Beger, Florian Walter, Alois Knoll
Chair of Robotics, Artificial Intelligence and Real-time Systems Technical University of Munich, Germany