Winner-Take-All Spiking Transformer for Language Modeling

Abstract: 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.

• The paper introduces a novelWTA-based attention mechanism that replaces softmax in spiking Transformers to reduce energy cost.
• It details two architectures, WE-Spikingformer and WD-Spikingformer, that deliver competitive accuracy on tasks like GLUE and QA while saving energy.
• Empirical evaluations demonstrate strong accuracy-energy trade-offs and scalability, highlighting the potential of spiking approaches for neuromorphic NLP.

Winner-Take-All Spiking Transformers for Energy-Efficient Language Modeling

Introduction

This work presents a systematic exploration of softmax-free, spike-driven Transformer architectures specifically tailored for language modeling tasks. By fusing the energy-efficient, event-driven computation of Spiking Neural Networks (SNNs) with the scalability of Transformers, the study proposes a biologically inspired Winner-Take-All (WTA) mechanism to address the prohibitively high energy costs and neuromorphic incompatibility of softmax-based self-attention. The approach yields two core architectures—WE-Spikingformer for masked language modeling and WD-Spikingformer for causal language modeling—built on novel WTA-based attention modules. Comprehensive evaluation across natural language understanding, question answering, and commonsense reasoning benchmarks demonstrates strong accuracy-energy trade-offs and positions the method as a promising candidate for general-purpose, neuromorphic LLMs.

Motivation: Spiking Transformers and the Energy Bottleneck

SNNs are attractive for neuromorphic computing due to their binary, event-driven processing that replaces traditional multiply-accumulate operations with sparse accumulates, delivering substantial theoretical energy savings. While spiking Transformers have achieved competitive results in vision tasks, a direct application to language modeling exacerbates key limitations:

• Existing architectures rely heavily on softmax-based attention, which is costly in terms of exponentiation, normalization, and floating-point operations—undermining SNNs' primary energy efficiency benefits.
• Non-spiking activations (such as GeLU or SiLU) persist in language-oriented SNN variants, further impeding adoption on neuromorphic hardware.
• Vision-oriented, softmax-free spiking Transformers (e.g., SpikFormer/SpikingFormer) transfer poorly to NLP: linguistic dependencies are longer and more global, and prior solutions evidence significant accuracy drop (e.g., 19.9% gap between SpikeBERT and BERT on GLUE).

To address these issues, the proposed WTA mechanism substitutes softmax, mimicking lateral inhibition in biological circuits to yield parsimonious, highly sparse attention maps—enabling efficient, fully spike-driven LLM pretraining and inference.

Winner-Take-All Mechanisms for Spiking Attention

Winner-Take-All selects and activates a single "winning" element (or top-K elements) from a set of competitors per position, enforcing extreme sparsity. In this context, it is used to replace softmax for self-attention in SNN Transformers.

• Hard WTA: Outputs a one-hot indicator for the maximum response—functionally, this is a $T\to0$ limit of temperature softmax.
• Top-K WTA: Activates the top $K$ scoring entries per query, generalizing Hard WTA for trade-offs between sparsity and expressivity.
• Sparsemax: Yields adaptive sparsity by projecting onto the probability simplex, facilitating differentiable, sparse attention.

All variants remove the need for exponentiation and costly normalization, and are intrinsically compatible with SNNs. Training leverages surrogate gradients: softmax gradients are used to propagate errors through the non-differentiable WTA selection during backpropagation, ensuring end-to-end differentiability.

Figure 1: The architecture of WTA Spiking Self-Attention (WSSA) and Causal WTA Spiking Self-Attention (CWSSA), contrasting the proposed softmax-free spike-driven approach with previous softmax-based SpikeLM attention.

WE-Spikingformer and WD-Spikingformer Architectures

Two primary model variants are introduced:

• WE-Spikingformer: An encoder-only, spike-driven Transformer for masked language modeling, built upon WTA Spiking Self-Attention (WSSA).
• WD-Spikingformer: A decoder-only, autoregressive Transformer for causal language modeling, utilizing Causal WTA Spiking Self-Attention (CWSSA) with a causal mask.

Both abandon non-spiking components in favor of exclusive use of spiking neurons (preferably NI-LIF for speed or T-LIF for expressivity), enabling pure spike-driven SNN training and inference. The multi-layer perceptron blocks in both models operate without floating-point nonlinearities, further optimizing for deployment on neuromorphic substrates.

Figure 2: Overview of the WE-Spikingformer (encoder-only, left) and WD-Spikingformer (decoder-only, right) architectures for spike-based language modeling.

Empirical Results: Accuracy and Efficiency on Language Understanding Benchmarks

The paper presents extensive evaluation across three task classes:

• Natural Language Understanding (GLUE): WE-Spikingformer outperforms preceding spike-driven softmax-free Transformers by a substantial margin—e.g., achieving 66.3% average accuracy vs. 59.7% for SpikeBERT and 34.6% for LIF-BERT at equal model scale. The model reduces the accuracy gap to softmax-based spiking competitors (SpikeLM), despite improved energy profiles.
• Question Answering and Commonsense Reasoning: WD-Spikingformer remains competitive with larger, energy-hungry baselines (e.g., Qwen-1.5B, SpikeLLM-7B). At only 0.4B parameters, it achieves 28.4% (QAT) and 43.2% (CRT) accuracy with approximately 7% (238.4 mJ vs 3398.3 mJ) of the energy budget.
• Scaling Analysis: Performance continues to improve with scaling: WE-Spikingformer-1.0B shows steady gains as parameter count or pretraining tokens are increased.

Figure 3: Performance scaling of WE-Spikingformer with increased model parameters for QA and commonsense reasoning.

Figure 4: Performance scaling of WE-Spikingformer-1.0B with increased pretraining tokens, demonstrating effective utilization of additional data for both QAT and CRT.

Ablations: WTA Effectiveness, Model Variants, and Surrogates

Multiple ablation studies validate the design choices:

• HardWTA vs. Softmax: Replacing softmax with HardWTA in self-attention yields minimal loss (<0.5% absolute) in accuracy on GLUE, QAT, or CRT, but with sharply reduced energy estimates and hardware complexity.
• WTA Variants: Top-K WTA and Sparsemax provide similar GLUE scores (≤1% variance) but HardWTA is preferred for maximal sparsity.
• Time Steps: Increasing SNN simulation steps further boosts accuracy, indicating trade-offs between latency, compute, and task performance.
• Surrogate Gradients: Softmax gradient outperforms alternatives (such as ReLU) as a surrogate for WTA in all evaluated tasks.

Practical and Theoretical Implications

The integration of WTA into spike-driven Transformer attention layers directly addresses the primary obstacles in scaling SNNs to NLP:

• By eliminating exponentiation and floating-point operations, the approach aligns with neuromorphic hardware constraints, enabling more realistic energy efficiency gains.
• The study empirically demonstrates that strict sparsity (as imposed by WTA) sustains competitive performance in challenging, non-vision domains, contradicting prior assumptions about the necessity of dense softmax normalization for self-attention.
• Parameter and data scaling analyses indicate an untapped capacity for large-scale SNN-based language modeling, suggesting emergent behaviors similar to conventional deep learning models as resources are increased.

The main caveat is the reliance on theoretical rather than measured energy estimates, and current models remain insufficient for very long-context, generative downstream tasks. Addressing neuromorphic deployment and further optimizing SNN processing for extensive context dependencies are compelling directions for future research.

Conclusion

This paper advances spike-driven language modeling by eliminating softmax and non-spiking nonlinearities from Transformer attention, replacing them with biologically-motivated, highly sparse Winner-Take-All mechanisms. The resulting WE-Spikingformer and WD-Spikingformer architectures achieve a new state-of-the-art for spike-driven Transformers across 16 NLP tasks, delivering substantial energy savings with negligible performance loss relative to dense baselines. The work substantiates the suitability of winner-take-all competition as a practical and theoretically justified alternative to softmax in neuromorphic NLP, and motivates further scaling and hardware implementation studies.

(2604.11321)