Source-linked AI summary
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
TL;DR
Compressed query, key, and value representations can create cross-time attention interactions absent from spike-driven inference, producing Temporal Interaction Mismatch. PT-SSA computes attention between reconstructed time-aligned slices in parallel, and the reported ImageNet-1K integer–spike Top-1 gap falls from 27.78 to 0.06 percentage points.
Problem
Applying E-SDSA to SFA query, key, and value activations introduces cross-time interactions absent from time-aligned spike inference, and ImageNet-1K spike inference loses more than 27 Top-1 points relative to integer-valued inference.
Source: §1 · p. 1Method
PT-SSA reconstructs virtual spike slices from I-LIF counts or SFA firing rates, computes attention only between time-aligned slices in parallel, and sums per-step outputs; Adaptive PT-SSA adds learned per-block rescaling before the SFA neuron.
Source: §1 · p. 2Results
Across CIFAR-10, CIFAR-100, and ImageNet-1K, the proposed methods reduce integer–spike gaps; on ImageNet-1K the Top-1 gap falls from 27.78 to 0.06 percentage points, with 74.53% spike-driven Top-1 accuracy.
Source: §1 · p. 2, §6 · p. 11Takeaways & Limitations
Under the corresponding I-LIF/SFA conversion assumptions, time-aligned attention preserves local train–inference equivalence through the subsequent spiking neuron and linear projection.
Source: §1 · p. 2, §4.2.3 · p. 5The stated local equivalence depends on the corresponding I-LIF/SFA conversion assumptions and extends through the subsequent spiking neuron and linear projection.
Source: §4.2.3 · p. 5
Abstract
from arXiv · showhide
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.
1. Introduction
Compressed query, key, and value representations can introduce cross-time interactions absent from spike-driven inference, a discrepancy termed Temporal Interaction Mismatch (TIM). The paper proposes time-aligned attention and adaptive rescaling, reporting a reduced integer–spike accuracy gap across vision benchmarks.
1. Introduction
- Applying E-SDSA to SFA activations introduces cross-time interactions during training, whereas temporal spike inference retains only time-aligned interactions.Reproducing cross-time interactions at inference would require additional computation and storage on resource-constrained neuromorphic hardware.Source: §1 · pp. 1–2
- PT-SSA reconstructs virtual spike slices from I-LIF counts or SFA firing rates, computes attention between time-aligned slices in parallel, and sums per-step outputs.Under SFA conversion assumptions, this alignment preserves local equivalence through the subsequent spiking neuron and linear projection.Source: §1 · p. 2
- Adaptive PT-SSA learns a positive per-block rescaling before the output SFA neuron to improve firing-level utilization.The rescaling accommodates the reduced attention-output scale under SFA.Source: §1 · p. 2
- 27.78 to 0.06 percentage points: Adaptive PT-SSA reduces the ImageNet-1K integer–spike Top-1 gap.The proposed methods also reduce gaps on CIFAR-10 and CIFAR-100.Source: §1 · p. 2
- The Triton-fused PT-SSA training kernel has comparable peak memory and 5.4% lower throughput than compressed integer SSA.The introduction reports that it retains efficiency close to compressed integer SSA.Source: §1 · p. 2
2. Related Work
Related work addresses temporal training cost through parallel computation or compact spike representations, alongside advances in spiking attention and adaptive quantization. The paper positions PT-SSA as addressing temporal alignment in SFA-based attention, with rescaling to accommodate its output scale.
2.1. Efficient SNN Training and Conversion
- Parallel spiking neurons and fixed-point methods reduce temporal recurrence, while I-LIF counts and SFA firing rates avoid explicit virtual-step simulation during training.The compact activations are expanded into temporal spike sequences at inference.Source: §2.1 · p. 2
2.2. Spiking Transformers and Attention
- Spiking Transformer work develops spike-based attention and revised backbones [Zhou et al., 2023b; Zhou et al., 2023a; Yao et al., 2023, 2024], while other work studies hierarchical attention and quantized attention.Conversion methods instead address nonlinear attention approximation or exact ANN–SNN equivalence.Source: §2.2 · p. 2
2.3. Adaptive Neurons and Quantization
- Adaptive PT-SSA rescales before each attention-output SFA neuron to accommodate the output-scale change caused by excluding cross-time terms.In the two-step example, compressed counts yield 4 while same-step products yield 2 despite identical total spike counts.Source: §2.3 · p. 3, Figure 1 · p. 3
3. Preliminaries
The preliminaries establish the neuron representations and spike-driven attention operators used to analyze the paper's train–inference discrepancy. They describe LIF dynamics, I-LIF and SFA conversions, and SDT-V3's E-SDSA.
3.1. LIF, I-LIF, and SFA
- LIF uses charging, firing, and reset dynamics whose membrane-state recurrence requires sequential simulation; I-LIF uses counts and SFA uses normalized firing rates during training.Counts are expanded into D-step spike sequences at inference, yielding TD total inference steps for T task-level steps.Source: §3.1 · p. 3
3.2. Spiking Self-Attention
- SSA performs softmax-free spatial attention separately at each time step, using spike-valued queries, keys, and values for sparse dot products and accumulation.Omitting softmax avoids dense floating-point normalization.Source: §3.2 · p. 4
3.3. Spike-Driven Transformer V3
- SDT-V3 combines SFA with E-SDSA, but the preliminaries note that SFA's linear-operator equivalence analysis does not establish equivalence through E-SDSA.At training, E-SDSA receives firing-rate tensors; inference converts counts to virtual spike sequences.Source: §3.3 · p. 4
4. Method
The method first analyzes the inconsistency from applying SDT-V3's E-SDSA to SFA firing-rate activations, then introduces PT-SSA and Adaptive PT-SSA. The analysis treats SFA as a normalized special case of I-LIF.
4. Method
- PT-SSA restores time-aligned attention during parallel training, and Adaptive PT-SSA rescales attention output before the standard SFA neuron.The rescaling accommodates PT-SSA's output scale under the normalized SFA representation.Source: §4 · p. 4
- The analysis uses s = 1 for I-LIF and s = D for SFA under its normalized-special-case treatment.These settings are specified for the analysis using Eq. (5).Source: §4 · p. 4
4.1. Analysis of Temporal Interaction Mismatch
Applying attention to compressed firing-count representations includes cross-time interactions, whereas temporal inference retains only aligned interactions. This mismatch means firing-count conservation alone does not guarantee operator-level equivalence.
4.1. Analysis of Temporal Interaction Mismatch
- Applying attention directly to spike-count representations therefore differs from aggregating attention over the corresponding temporal spike slices.The discrepancy arises because a time-local trilinear attention operator is evaluated on temporally compressed representations.Source: §4.1 · p. 4
- Compressed attention includes all temporal index triplets, while inference retains only terms whose query, key, and value indices match.The difference consists of cross-time terms; when D=2, these include Q0K1 and Q1K0, and the paths coincide only in degenerate cases such as D=1 or vanishing off-diagonal interactions.Source: §4.1 · p. 4
4.2. Parallel Time-Aligned Spiking Self-Attention
PT-SSA reconstructs virtual spike slices from compact query, key, and value representations and computes attention only between slices at matching virtual steps. Its parallel outputs are summed, preserving local train–inference equivalence under the stated conversion assumptions.
4.2. Parallel Time-Aligned Spiking Self-Attention
- PT-SSA reconstructs spike sequences and evaluates attention only within matching virtual steps, then sums the per-step outputs.The shared virtual-step index excludes cross-time products; the block applies the factor α/s^3 after summation.Source: §4.2 · p. 4, §4.2.2 · pp. 4–5, Figure 2 · p. 5
4.2.1. Spike Sequence Reconstruction
- Each reconstructed spike slice is determined by the compact representation and virtual-step index, so all D slices can be generated in parallel.For an entry with firing count c, the reconstruction emits a spike exactly when d < c.Source: §4.2.1 · p. 4
4.2.2. Parallel Time-Aligned Spiking Self-Attention
- At each virtual step, independent attention outputs can be computed in parallel and reduced after the attention operations.Only query, key, and value spikes with the same virtual-step index interact.Source: §4.2.2 · pp. 4–5
4.2.3. Time-Alignment Property
- PT-SSA and spike-driven inference use the same virtual-step spike slices and attention outputs, yielding identical attention outputs before firing.Under the corresponding I-LIF/SFA conversion assumptions and rules, local equivalence extends through the subsequent spiking neuron and linear projection.Source: §4.2.3 · p. 5
4.2.4. Parallel Surrogate-Gradient Routing
- A unit straight-through estimator routes gradients through spike reconstruction, and the D virtual-step gradients are computed in parallel before reduction.The normalized-slice implementation gives the same forward output and surrogate gradient as the stated equations.Source: §4.2.4 · pp. 5–6
4.3. Adaptive Rescaling for SFA
PT-SSA removes nonnegative cross-time terms, which can lower the membrane potential reaching the SFA output neuron and leave firing-count levels underused. Adaptive PT-SSA adds a learnable positive rescaling before that neuron.
4.3. Adaptive Rescaling for SFA
- Adaptive PT-SSA adds a learnable positive rescaling before each attention-output SFA neuron to address potentially low firing counts after cross-time terms are removed.The block-shared rescaling adds L scalar parameters for a network with L blocks and is equivalent to adapting the effective firing threshold relative to the unscaled output.Source: §4.3 · p. 6
4.3.1. Rescaling and Firing
- The rescaling divides the PT-SSA output by θ_l before standard unit-threshold SFA firing, with θ_l=1 giving the identity operation.At inference, the accumulated output can use the equivalent firing threshold s^3θ_l/α without explicit scaling and division.Source: §4.3.1 · p. 6
4.3.2. Learning the Rescaling Parameter
- The learned parameter is represented by a clipped exponential initialized from a_l=0, and an identity STE omits the 1/θ_l factor in the gradient routed to PT-SSA.The effect of this gradient-routing choice is evaluated in the component-wise ablation in Section 5.5.4.Source: §4.3.2 · p. 6
4.4. Efficient Implementation
A Triton-fused PT-SSA kernel reduces training memory overhead by reconstructing slices within tiles and accumulating their attention outputs without materializing all intermediates.
4.4. Efficient Implementation
- The Triton implementation reconstructs slices and computes parallel time-aligned attention, avoiding simultaneous storage of all reconstructed tensors and attention intermediates.During backpropagation, it regenerates required slices and intermediates, trading recomputation for lower training memory.Source: §4.4 · pp. 6–7
5. Experiments
Experiments evaluate operator-level alignment, classification accuracy, sensitivity to virtual expansion length, training efficiency, and output-neuron firing response on CIFAR-10, CIFAR-100, and ImageNet-1K. PT-SSA variants reduce train–inference gaps across the tested datasets and representations.
5.1.1. Datasets and Backbones
- The evaluation covers CIFAR-10, CIFAR-100, and ImageNet-1K, using CIFAR-adapted E-SpikeFormer-4-384 and the 5.11M-parameter E-SpikeFormer-8-192 backbone for ImageNet.Within controlled comparisons, SSA and PT-SSA share the backbone topology, projections, convolutional blocks, residual paths, classifier, and attention hyperparameters.Source: §5.1.1 · p. 7
5.1.2. Training Protocol
- CIFAR models use dataset-specific Spikingformer recipes, while ImageNet models follow Spike-Driven Transformer V3 recipes and train for 200 epochs on eight GPUs.The protocol specifies optimizer, augmentation, and schedule settings separately by dataset.Source: §5.1.2 · p. 7
5.1.3. Representations and Evaluation
- Inference paths use the same checkpoint; at spike-driven inference, each count is expanded into a D-step spike sequence and attention is evaluated over D virtual steps.CIFAR results report means and sample standard deviations over two seeds; ImageNet uses one seed.Source: §5.1.3 · p. 7
5.2. Operator-Level TIM
- At D=4, PT-SSA reduces relative L2 error from 0.7136 to 2.11×10^-4 for I-LIF and from 7.4120 to 1.60 × 10^-4 for SFA.Cosine similarity rounds to 1.000 in both cases; the comparison isolates attention computation from preceding or subsequent layers.Source: §5.2 · p. 8
5.3.2. SFA Results
- 0.06 percentage points: Adaptive PT-SSA reduces the ImageNet-1K gap from 27.78 percentage points while reaching 74.53% spike-driven Top-1 accuracy.Its integer-valued accuracy is 74.59%, 0.81 percentage points below SSA, so the recovered spike-driven accuracy does not reflect a stronger integer-valued model.Source: Figure 4 · p. 8, §5.3.2 · p. 9
5.3.1. I-LIF Count Results
- PT-SSA reduces the CIFAR-100 mean gap from 1.85 ± 0.08 to 0.29 ± 0.04 percentage points; on CIFAR-10, it falls from 0.41 ± 0.27 to 0.25 ± 0.04 percentage points.On CIFAR-100, mean spike-driven inference accuracy rises from 78.84 ± 0.03% to 80.10 ± 0.18%, while mean integer-valued inference accuracy decreases from 80.69 ± 0.11% to 80.39 ± 0.14%.Source: §5.3.1 · p. 8
5.3.2. SFA Results
- 5.58 to 0.05 percentage points: Adaptive PT-SSA reduces the CIFAR-10 SFA mean gap, and on CIFAR-100 reduces it from 26.99 to 0.07 percentage points.CIFAR-100 spike-driven accuracy rises from 53.83 ± 1.53% to 82.42 ± 0.46%.Source: Figure 4 · p. 8, §5.3.2 · p. 9
5.4. Training Efficiency
PT-SSA’s training efficiency and ablations show that time-aligned attention limits train–inference gaps as virtual expansion grows, while adaptive thresholds improve SFA firing-level utilization. The Triton implementation retains a throughput advantage over recurrent LIF SSA.
5.4. Training Efficiency
- PT-SSA (Triton) reaches 2.91× recurrent LIF SSA throughput while reducing peak memory by 84.1%, and runs only 5.4% below I-LIF SSA throughput.Relative to PT-SSA (Torch), Triton improves throughput by 15.3% and reduces peak memory by 12.5%.Source: §5.4 · p. 9
5.5.1. Representation and Virtual Expansion Length
- As D increases, SSA includes more off-diagonal triplets while PT-SSA retains only D time-aligned triplets; under SFA, D also changes membrane-potential scale relative to the output neuron’s firing threshold.Figure 5 plots mean absolute Top-1 train–inference gap against D on CIFAR-100, with separate I-LIF and SFA panels.Source: §5.5.1 · p. 9, Figure 5 · p. 10
5.5.2. Effect of 𝐷on TIM
- Under I-LIF, SSA’s mean gap rises from 0.62 to 4.34 percentage points as D grows from 2 to 8, while PT-SSA remains at or below 0.29.At D=8, PT-SSA improves mean spike-driven accuracy by 3.31 percentage points despite 0.81 percentage points lower mean integer-valued accuracy.Source: §5.5.2 · p. 9
- Under SFA, SSA’s mean gap grows from 0 at D=1 to 7.54 percentage points at D=2 and above 25 at D=4 and 8; Adaptive PT-SSA stays at or below 0.28.At D=8, its gap is larger than at D=4.Source: §5.5.2 · pp. 9–10
5.5.3. PT-SSA Component Ablation
- At D=4, PT-SSA raises SFA spike-driven accuracy from 53.83 ± 1.53% to 81.03 ± 0.01% and reduces the mean gap from 26.99 to 0.25 percentage points.Learning the firing threshold further improves mean integer-valued and spike-driven accuracy by 1.58 and 1.39 percentage points, respectively.Source: §5.5.3 · p. 10, Table 4 · p. 10
5.5.4. Gradient Routing for the Adaptive Threshold
- Decoupled threshold gradients improve integer-valued and spike-driven Top-1 accuracy and reduce the mean train–inference gap to 0.07 ± 0.04 percentage points.Across two runs, integer-valued Top-1 rises from 81.05 ± 0.46% to 82.35 ± 0.42%, and spike-driven Top-1 rises from 81.46 ± 0.18% to 82.42 ± 0.46%; the coupled-versus-decoupled comparison does not isolate the shared normalization factor’s effect.Source: §5.5.4 · p. 10, Table 4 · p. 10
5.6. Analysis of the Adaptive Firing Threshold
- The fixed 1/D^3 threshold yields a 97.73% firing rate but only 0.367 bits of entropy, whereas decoupled adaptive thresholds yield 17.97% and 1.519 bits.The adaptive setting also attains the highest spike-driven accuracy among the compared settings; a higher mean firing rate alone does not ensure higher accuracy.Source: §5.6 · p. 10
6. Conclusion
The paper identifies Temporal Interaction Mismatch between compressed integer-valued training and temporal spike-driven inference in SFA-based self-attention. PT-SSA uses time-aligned virtual spike slices, and experiments report smaller accuracy gaps alongside efficient training.
6. Conclusion
- Temporal Interaction Mismatch is a structural inconsistency between compressed integer-valued training and temporal spike-driven inference in SFA-based spiking self-attention.The work identifies this as an operator-level issue.Source: §6 · p. 11
- PT-SSA reconstructs consecutive virtual spike slices and attends only across time-aligned query, key, and value slices; Adaptive PT-SSA learns per-block firing thresholds for SFA outputs.The approach sums per-step outputs as described in the paper context.Source: §6 · p. 11
- On ImageNet-1K, the Top-1 gap falls from 27.78 to 0.06 percentage points while spike-driven accuracy reaches 74.53%; the Triton-fused kernel delivers 2.91× recurrent LIF SSA throughput.Experiments also cover CIFAR-10 and CIFAR-100.Source: §6 · p. 11