Uncertainty-Aware Token Importance Estimation in Spiking Transformers

Technical Summary: Uncertainty-Aware Token Importance Estimation in Spiking Transformers

Problem Statement
Spiking Transformers have demonstrated significant potential for neuromorphic vision by combining the energy efficiency of Spiking Neural Networks (SNNs) with the global modeling capabilities of self-attention. However, a critical bottleneck remains: the multi-step token processing inherent to SNNs often introduces substantial redundancy and inference costs. While existing token reduction methods (e.g., token pruning or merging) aim to improve efficiency, they predominantly rely on response-driven cues such as activation magnitude, firing statistics, or feature similarity.

The authors argue that these criteria are insufficient for Spiking Transformers because they fail to explicitly characterize how token-level class evidence evolves over time. Unlike static transformers where token representations are determined at a single instant, spiking representations are progressively formed across multiple simulation steps. Consequently, token importance cannot be adequately inferred from single-step responses alone; it requires a temporal perspective that captures the evolution of evidence and uncertainty throughout the spiking dynamics.

Methodology: Uncert
To address this, the authors propose Uncert, a training-free, plug-and-play framework for estimating token importance based on temporal uncertainty patterns. The methodology consists of four main components:

1. Spiking Token Hierarchy: The framework utilizes a hierarchical spiking transformer backbone (e.g., QKFormer or MaxFormer). Token representations evolve across $T$ time steps, forming a temporal trajectory for each spatial token.
2. Token Uncertainty Criterion: Instead of using raw activation values, Uncert models token-wise class evidence using a Dirichlet distribution.
   • Classifier logits for a token $i$ at time $t$ are converted into non-negative evidence values via a Softplus activation.
   • These evidence values define the parameters ($\alpha$) of a Dirichlet distribution.
   • Token uncertainty ($U$) is derived from the total evidence ($S$) as $U = C/S$, where $C$ is the number of classes. This yields a temporal uncertainty trajectory $\{U^t_i\}_{t=1}^T$ for each token.
3. Spiking-Temporal Uncertainty Modeling: The framework summarizes the temporal uncertainty trajectory of each token using two statistics:
   • Mean ($\mu$): Represents the overall ambiguity level of the token.
   • Standard Deviation ($\sigma$): Captures the temporal fluctuation or stability of the token's uncertainty.
   • The final importance score is computed as $S_i = \mu_i + \lambda \sigma_i$, where $\lambda$ is a balancing coefficient (empirically set to 0.9). The authors observe that tokens with higher uncertainty or stronger fluctuations often correspond to informative regions where the model is resolving ambiguity, whereas low-uncertainty tokens often represent redundant background.
4. Token Selection: The calculated scores guide inference-time token reduction via two strategies:
   • Token Pruning: Selecting the top-$K$ tokens with the highest scores to retain, discarding the rest before self-attention computation.
   • Token Merging: Aggregating low-importance tokens into nearby high-scoring anchors.

Key Contributions

• Temporal Perspective on Token Importance: The paper revisits token importance estimation in Spiking Transformers, demonstrating that tokens exhibit heterogeneous uncertainty trajectories (stable, fluctuating, or persistently ambiguous) rather than uniform dynamics.
• Uncert Framework: The proposal of a training-free, plug-and-play method that models token evidence with a Dirichlet distribution and derives importance scores from temporally aggregated uncertainty statistics (mean and fluctuation).
• Empirical Validation: Extensive experiments on both static (CIFAR-10, CIFAR-100) and neuromorphic (DVS-CIFAR-10, DVS-Gesture) benchmarks validate the method's effectiveness.

Experimental Results

• Accuracy-Efficiency Trade-off: Uncert achieves favorable trade-offs between accuracy and efficiency. On CIFAR-100 and DVS-CIFAR-10, retaining only 80–90% of tokens often results in accuracy comparable to or slightly better than the full baseline.
• Pruning vs. Merging: Token pruning emerges as the most effective application of the uncertainty criterion. For instance, on CIFAR-100 with an 80% retention ratio, Uncert (Prune) achieved 81.06% accuracy, outperforming the baseline (80.98%) and random selection (80.40%). In contrast, token merging showed less consistent gains and often underperformed the baseline under aggressive compression.
• Component Analysis: Ablation studies confirm that combining both the mean uncertainty and temporal fluctuation yields superior performance compared to using either metric alone. This highlights that both the overall ambiguity level and the temporal stability of tokens are critical for distinguishing informative signals.
• Efficiency Gains: Block-level operations (SOP) were reduced by approximately 14.86% when the keep ratio was lowered to 0.40, demonstrating measurable computational savings.

Significance and Claims
The paper claims that Uncert offers a new perspective on token dynamics in Spiking Transformers by shifting the focus from instantaneous response strength to temporal uncertainty patterns. The authors posit that uncertainty is not merely a measure of instability but a useful signal for identifying discriminative content that requires further resolution.

The significance of the work lies in its ability to improve inference efficiency without requiring additional training or learnable predictors. By leveraging the intrinsic temporal evolution of spiking representations, Uncert provides a principled, training-free mechanism to reduce redundancy. The authors conclude that temporal uncertainty statistics offer a robust and effective cue for token selection, particularly for pruning, enabling more efficient Spiking Transformers for both static and event-based vision tasks.