Parallel Time-Aligned Spiking Self-Attention for Consistent Integer-Valued Training and Spike-Driven Inference

Peng Xue; Wei Fang; Kaiwei Che; Qingyan Meng; Zhengyu Ma; Yonghong Tian; Huihui Zhou

cs.NEOct 2, 2026

Parallel Time-Aligned Spiking Self-Attention for Consistent Integer-Valued Training and Spike-Driven Inference

Authors: Peng Xue, Wei Fang, Kaiwei Che, Qingyan Meng, Zhengyu Ma, Yonghong Tian, Huihui Zhou

Organizations: Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences, China · Pengcheng Laboratory, China · University of Chinese Academy of Sciences, China · Shenzhen Graduate School, Peking University, China ·

Abstract

Integer-valued leaky integrate-and-fire (I-LIF) neurons and spike firing approximation (SFA) reduce temporal training cost by representing spike trains as firing counts and normalized firing rates, respectively. However, applying spiking self-attention (SSA) directly to these compressed query, key, and value representations introduces cross-time interactions that are absent during spike-driven inference. We term this operator-level discrepancy Temporal Interaction Mismatch (TIM). We propose Parallel Time-Aligned Spiking Self-Attention (PT-SSA), which reconstructs consecutive virtual spike slices from either I-LIF counts or SFA firing rates, computes attention only between time-aligned slices in parallel, and sums the per-step outputs. To accommodate the reduced attention output scale under SFA, we further introduce Adaptive PT-SSA, which learns a positive per-block rescaling before the output SFA neuron to improve firing-level utilization. Experiments on CIFAR-10, CIFAR-100, and ImageNet-1K show that the proposed methods substantially reduce train--inference mismatch. On CIFAR-100 with I-LIF, PT-SSA reduces the mean Top-1 gap from 1.85 to 0.29 percentage points. On ImageNet-1K, Adaptive PT-SSA reduces the Top-1 gap from 27.78 to 0.06 percentage points and achieves 74.53% spike-driven Top-1 accuracy. A Triton-fused PT-SSA training kernel retains a $2.91\times$ throughput advantage over recurrent LIF SSA.

Figures & tables

Figure 1: Temporal interaction mismatch in a two-step I-LIF example (

D=2

s=1

). (a) Multiplying compressed spike counts includes cross-time products and yields the highlighted value 4. (b) Summing same-step spike products yields 2 despite the same total spike counts.

Figure 2: PT-SSA within a spiking attention block. (a) Compact queries, keys, and values enter PT-SSA, followed by output rescaling and firing. (b) Reconstructed spike slices are grouped by virtual step, and their time-aligned attention outputs are summed over the virtual time-step dimension. Red spike symbols indicate I-LIF or SFA spiking neurons, and

\times

indicates matrix multiplication. The illustrated SFA path uses

s=D

and a learned

\theta_{l}

; I-LIF uses

s=1

and

\theta_{l}=1

Figure 3: Adaptive rescaling and gradient routing for SFA-based PT-SSA. (a) Forward rescaling uses a learned

\theta_{l}

, while an identity STE passes the gradient from

\mathbf{U}_{l}

to the attention output unchanged. (b) Quantized firing count versus unscaled attention output at

D=4

for

\theta_{l}=1

and

0.2

; these curves show rounded SFA responses rather than temporal IF threshold crossings. (c) Local input-gradient gain versus

\theta_{l}

for ordinary (

1/\theta_{l}

) and decoupled (

1

) differentiation.

Dataset	Backbone	Representation	Attention path	Integer Top-1 (%) $\uparrow$	Spike-driven Top-1 (%) $\uparrow$	$\|\Delta\|$ (pp) $\downarrow$
CIFAR-10	E-SpikeFormer-4-384	I-LIF count	SSA	$\mathbf{97.26\pm 0.06}$	$96.85\pm 0.33$	$0.41\pm 0.27$
		I-LIF count	PT-SSA	$97.20\pm 0.07$	$\mathbf{96.95\pm 0.11}$	$\mathbf{0.25\pm 0.04}$
		SFA rate	SSA	$97.08\pm 0.19$	$91.50\pm 0.00$	$5.58\pm 0.19$
		SFA rate	Adaptive PT-SSA	$\mathbf{97.23\pm 0.03}$	$\mathbf{97.23\pm 0.09}$	$\mathbf{0.05\pm 0.01}$
CIFAR-100	E-SpikeFormer-4-384	I-LIF count	SSA	$\mathbf{80.69\pm 0.11}$	$78.84\pm 0.03$	$1.85\pm 0.08$
		I-LIF count	PT-SSA	$80.39\pm 0.14$	$\mathbf{80.10\pm 0.18}$	$\mathbf{0.29\pm 0.04}$

Table 1: Train–inference alignment for I-LIF counts and SFA firing rates at

D=4

. Integer-valued and spike-driven results in each row come from the same checkpoint. The paired absolute Top-1 gap

|\Delta|

is reported in percentage points and computed before rounding.

Neuron	Attention implementation	Peak memory (GiB) $\downarrow$	Forward (ms) $\downarrow$	Backward (ms) $\downarrow$	Throughput (sample/s) $\uparrow$
LIF	SSA	4.193	12.909	19.197	1993.4
I-LIF	SSA	0.679	4.099	6.348	6126.5
I-LIF	PT-SSA (Torch)	0.761	5.211	7.517	5028.3
I-LIF	PT-SSA (Triton)	0.666	4.522	6.518	5797.1

Table 2: Training efficiency on CIFAR-100 Spikingformer using one RTX 5090 (

B=64

D=4

, AMP FP16). SSA uses Torch, while recurrent LIF neurons use the SpikingJelly CuPy backend. Timings are averaged over 50 synchronized iterations after 20 warmup iterations and include forward and backward computation; data loading, optimizer updates, and Triton compilation are excluded. Peak memory excludes optimizer state; bold and underlined values indicate the best and second-best results.

Figure 4: Operator-level and end-to-end train–inference alignment at

D=4

(logarithmic horizontal axes; lower is better). (a) Relative

L_{2}

error against the aligned spike reference on CIFAR-100 (I-LIF) and ImageNet-1K (SFA). (b) Absolute Top-1 gap in percentage points on CIFAR-10 (C10), CIFAR-100 (C100), and ImageNet-1K (IN-1K).

Representation	Attention Impl.	$D$	Integer Top-1 (%) $\uparrow$	Spike-driven Top-1 (%) $\uparrow$	$\|\Delta\|$ (pp) $\downarrow$
I-LIF count	SSA	1	$77.31\pm 0.42$	$77.31\pm 0.42$	$0.00\pm 0.00$
		2	$79.88\pm 0.16$	$79.26\pm 0.04$	$0.62\pm 0.12$
		4	$80.69\pm 0.11$	$78.84\pm 0.03$	$1.85\pm 0.08$
		8	$80.90\pm 0.14$	$76.57\pm 0.39$	$4.34\pm 0.53$
	PT-SSA	1	$77.31\pm 0.42$	$77.31\pm 0.42$	$0.00\pm 0.00$
		2	$80.12\pm 0.11$	$79.91\pm 0.12$	$0.22\pm 0.01$

Table 3: CIFAR-100 accuracy and train–inference gap across virtual expansion lengths. Values are means

\pm

sample standard deviations over two runs.

Figure 5: Mean absolute Top-1 train–inference gap versus virtual expansion length

D

on CIFAR-100. Panels (a) and (b) show I-LIF and SFA, respectively; their vertical axes report percentage points on the same symmetric-log scale with a linear region around zero.

Setting	$\theta_{l}$ (mean $\pm$ std)	Firing rate (%)	Entropy (bits) $\uparrow$	Integer Top-1 (%) $\uparrow$	Spike-driven Top-1 (%) $\uparrow$	$\|\Delta\|$ (pp) $\downarrow$
Fixed unit	$1.00\pm 0.00$	10.94	1.204	$80.78\pm 0.09$	$81.03\pm 0.01$	$0.25\pm 0.10$
Fixed $1/D^{3}$	$0.02\pm 0.00$	97.73	0.367	$80.48\pm 0.40$	$80.61\pm 0.40$	$0.13\pm 0.01$
Adaptive (coupled)	$0.21\pm 0.02$	16.67	1.487	$81.05\pm 0.46$	$81.46\pm 0.18$	$0.42\pm 0.28$
Adaptive (decoupled)	$0.20\pm 0.03$	17.97	1.519	$\mathbf{82.35\pm 0.42}$	$\mathbf{82.42\pm 0.46}$	$\mathbf{0.07\pm 0.04}$

Table 4: Adaptive-threshold and gradient-routing analysis for SFA-based PT-SSA on CIFAR-100 (

D=4

). Top-1 and gap entries are means

\pm

sample standard deviations over two runs. Firing rates and entropies are computed on the complete test set for each block and run, then averaged;

\theta_{l}

statistics summarize four blocks and two runs.

Spiking Neural Networks

Spiking neural networks process information through discrete, time-dependent neural events, aiming for efficient computation and biologically inspired learning.

See all 336 papers

Lapis: Laplacian Spiking Attention via First-Spike Timing and Membrane Leakage
Aug 12, 2026Kaiwen Tang, Jiaqi Zheng, Zixuan Zhu +3Spiking Transformers Efficient Attention
School of Computing, National University of Singapore · 2Sea AI Lab · 4Shanghai Advanced Research Institute, Chinese Academy of Sciences +2
FidLIF: Parallel Training of High-Fidelity Spiking Neurons with Exact BPTT
Oct 12, 2026Kaixiang He, Yujie WuLIF Neurons Spiking Neural Networks
Department of Computing, The Hong Kong Polytechnic University
Scalable Training of Continuous-Time Spiking Neural Networks with Differentiable Spike-Time Discretization
Jul 16, 2026Yusuke Sakemi, Tomoya Takeuchi, Takeo Hosomi +1Spiking Neural Networks LIF Neurons
Research Center for Mathematical Engineering, Chiba Institute of Technology, Narashino, Japan · International Research Center for Neurointelligence (WPI-IRCN), The University of Tokyo, Tokyo, Japan · NEC Corporation, Kawasaki, Japan

Aug 12, 2026J/K move · Enter open · S save

Lapis: Laplacian Spiking Attention via First-Spike Timing and Membrane Leakage

Kaiwen Tang, Jiaqi Zheng, Zixuan Zhu, Yiqun Wang+2

School of Computing, National University of Singapore · 2Sea AI Lab · 4Shanghai Advanced Research Institute, Chinese Academy of Sciences+2

Spiking Transformers Efficient Attention Spiking Neural Networks Attention Mechanisms

Open paper AskPDF arXiv

Self-attention has become central to spiking vision transformers, yet its query-key scoring is still largely inherited from dense networks. Existing spiking variants either simplify dot product scoring or replace it with discrete operators, but spike timing, the native variable of a spiking network, does not directly define how tokens are related. We propose Lapis, a spiking attention mechanism that scores each token pair by the L1 distance between its query and key first-spike latency vectors under time-to-first-spike coding, and maps this distance to an affinity through a Laplacian kernel. The kernel's exponential decay matches the impulse response of a leaky integrate-and-fire membrane, so the accumulated latency difference determines the decay of a membrane trace, while row normalization reduces to a bit shift under power-of-two rounding. Scoring therefore needs only subtraction, absolute value, and accumulation, and removes all multiplication between query and key channels. Under a matched backbone and training schedule, Lapis reaches 96.56% top-1 accuracy on CIFAR-10, within 0.53 points of dot-product scoring. On ImageNet-1K, it reduces the estimated arithmetic energy of the attention path by 14.5x relative to dense dot-product attention. The deployed 6-bit model attains 83.25% top-1 accuracy at an estimated arithmetic energy of 3.28mJ per image.

Abstract

Spiking Neural Networks

Explore similar work

Lapis: Laplacian Spiking Attention via First-Spike Timing and Membrane Leakage

FidLIF: Parallel Training of High-Fidelity Spiking Neurons with Exact BPTT

Scalable Training of Continuous-Time Spiking Neural Networks with Differentiable Spike-Time Discretization