技術要約: SAGE – Attention-Guided Entropyによるサロゲート勾配の適応化を用いたSpiking Transformer

1. 問題提起

スパイキングニューラルネットワーク（SNN）は、イベント駆動型の疎な計算を利用することで、従来のディープニューラルネットワークに代わるエネルギー効率の高い選択肢を提供します。しかし、深層Spiking Vision Transformer（ViT）の学習は、スパイク生成関数（ヘヴィサイド階段関数）の非微分性により困難が伴います。このため、バックプロパゲーションにおいて微分の近似を行うために、**サロゲート勾配（surrogate gradients）**を使用する必要があります。

現在の最先端のアプローチは、通常、すべての空間トークンおよび訓練ステージに対して、固定された、あるいはグローバルに一様なサロゲート勾配関数（例：静的なアークタンジェントやFast Sigmoid）を採用しています。本論文は、このような空間的に一様な近似が、深層Spiking Transformerにおいて決定的な失敗モードを引き起こすと主張しています。

• トークンレベルのダイナミクスの無視: 視覚的トークンは、さまざまなレベルのセマンティックな不確実性を備えています。高信頼度のトークン（例：明確な物体の境界）には精密な勾配更新が必要ですが、高不確実性のトークン（例：曖昧な背景）は、より広範で探索的な勾配の流れを必要とします。
• 最適化の失敗: すべてのトークンに対して静的なサロゲートウィンドウを強制することは、文脈に応じたクレジット割り当て（credit assignment）を提供できず、失敗を招きます。深い層において、これは深刻な勾配消失を引き起こし、アテンションヘッドが「デッドヘッド（dead heads）」（膜電位が閾値以下に留まり、重みの更新が停滞する状態）へと崩壊する原因となります。

核心となる課題は、推論時に求められる決定論的かつ低電力なバイナリ実行を損なうことなく、局所的なトークンの不確実性に合わせて勾配の流れを動的に調整する、適応型サロゲートメカニズムを開発することです。

2. メソドロジー: SAGEフレームワーク

著者らは、アテンションエントロピーの分散に基づいてサロゲート勾配の傾きを調整する、不確実性を考慮したフレームワークであるSAGE（Surrogate-gradient Adaptation via Attention-Guided Entropy）を提案しています。

2.1 コアコンセプト

SAGEは、訓練フェーズにおいてのみ動作します。本手法は、自己アテンションヘッドにおけるアテンションエントロピーの分散を分析することで、各トランスフォーマーブロックの不確実性を推定します。この信号は、サロゲート勾配関数の傾き（$\alpha$）を適応的に調整するために使用されます。重要な点として、推論モデルは変更されません。不確実性コントローラーは訓練後に除去されます。

2.2 アテンション不確実性の推定

本フレームワークは、SpikformerアーキテクチャのMulti-Head Self-Attention (MHA) モジュールを活用します。

1. 生のアテンションスコア: 標準的なSpikformerは、Softmax正規化を行わずに、生のスコア $S = s QK^\top$ を計算します。
2. 補助的な分布: 不確実性推定のためのみに、SAGEは固定温度（$T_{ent} = 0.25$）を用いたSoftmaxとdetach操作を適用することで、一時的な確率分布 $A$ を構築します。
   $$A = \text{Softmax}\left(\frac{\text{detach}(S)}{T_{ent}}\right)$$
3. エントロピー計算: 各アテンションヘッド $i$ に対して、正規化されたシャノンエントロピー $E_i$ が計算されます。
4. 不確実性信号: 平均エントロピーではなく、SAGEは全 $H$ 個のヘッド間におけるエントロピーの**標準偏差（分散）**を計算します。
   $$D = \text{Std}(E_1, E_2, \dots, E_H)$$
   高い分散はヘッド間の不一致（高い不確実性）を示し、低い分散は一貫したアテンションパターンを意味します。

2.3 不確実性ガイド型サロゲート変調

統計量 $D$ は、ブロック固有のサロゲート傾き $\alpha$ を生成するために処理されます。

1. 平滑化: 安定性を確保するため、瞬時の分散は指数移動平均（EMA）を用いて平滑化されます。
2. 正規化: 平滑化された値は、実行中の平均および標準偏差の統計量を用いて正規化され、すべてのトランスフォーマーブロック間で中心化され、相対的な不確実性スコア $\tilde{z}$ を得ます。
3. 適応的マッピング: デッドゾーンコントローラーが微小な変動を抑制します。傾きは双曲線正接関数を介して変調されます。
   $$\hat{\alpha} = \begin{cases} 4.0 & \text{if } |\tilde{z}| < 0.25 \\ 4.0 + 0.5 \tanh(\tilde{z}) & \text{otherwise} \end{cases}$$
   最終的な傾きは区間 $[3, 5]$ にクリップされます。
4. 統合: この適応的な $\alpha$ は、バックプロパゲーション（逆伝播）の間のみ適用されます。フォワードパスおよび推論アーキテクチャは、ベースラインのSpikformerと同一です。

2.4 計算オーバーヘッド

追加される計算は、エントロピーの計算（$O(HN)$）と標準偏差の計算（$O(H)$）であり、これらはアテンションメカニズムの$O(N^2)$ の複雑さに比べれば無視できる程度です。プロファイリングによれば、訓練中のオーバーヘッドは約 0.03 ms/ミニバッチ であり、推論時のオーバーヘッドは ゼロ です。

3. 主な貢献

本論文は、主に3つの貢献を述べています。

1. 初の不確実性認識フレームワーク: アテンション由来の不確実性を利用して、スパイキングTransformer向けに設計された初の適応型サロゲート勾配フレームワークの導入。
2. エントロピー分散としての信号: ヘッド間のアテンションエントロピーの分散（標準偏差）が、サロゲート適応のための代替指標（Gini不純度や平均エントロピーなど）よりも、より識別力があり信頼性の高いオンライン不確実性信号であることを実証。
3. 性能と効率性: 固定および学習可能なベースラインと比較して、無視できる程度の訓練オーバーヘッドでありながら、一貫した性能向上を競争力のあるデータセットで検証。

4. 実験結果

実験は、Spikformer-4-384バックボーンを用い、CIFAR-10、CIFAR-100、および ImageNet-200 で行われました。

• 精度向上: SAGEは、CIFARデータセットにおいて固定サロゲートのベースラインに対して一貫した改善を達成しました。
  • CIFAR-10: Top-1精度 96.32%（ベースラインのSpikformerの95.51%に対し）。
  • CIFAR-100: Top-1精度 78.69%（ベースラインの78.21%に対し）。
  • ImageNet-200: このベンチマークにおいて、SAGEは60.60%のTop-1精度を達成しましたが、これはベースラインのSpikformerの70.24%と比較して低下しています。論文では、SAGEがCIFARデータセットで性能を向上させる一方で、ImageNet-200の結果はベースラインより低くなったことが記されていますが、同手法は低レイテンシ領域での有効性と、元のアーキテクチャを維持できる能力において強調されています。
• 時間的堅牢性: SAGEは、より低いシミュレーションタイムステップ（$T \le 4$）、特に標準的な設定である $T=4$ において最も顕著な改善を示しました。タイムステップが増加するにつれ（$T=8, 12$）、性能差は縮まり、この手法が低レイテンシアプリケーションに最も効果的であることを示唆しています。
• 比較: SAGEは、複数のデータセットおよびタイムステップにおいて、固定サロゲート勾配および「学習可能（Learnable）」な構成（単一のグローバルパラメータを訓練するもの）の両方を上回りました。

5. 重要性と主張

本論文は、SAGEを、深層Spiking Transformerの特定の限界に対処するプラグアンドプレイ型の最適化戦略として位置付けています。

• アーキテクチャの保存: 本手法は、元のネットワークアーキテクチャ、モデルパラメータ、および推論パイプラインを保持します。モデル自体に追加の学習可能パラメータを導入せず、推論レイテンシも発生させません。
• 文脈的クレジット割り当て: ローカルなトークンの不確実性に基づいて勾配の傾きを動的に調整することにより、SAGEは深い層における「デッドヘッド」現象と勾配消失を緩和し、ネットワークが訓練中に曖昧な特徴をより適切に扱えるようにします。
• 実用性: 著者らは、SAGEが軽量なソリューションであり、既存の訓練パイプラインにシームレスに統合できるため、展開時の変更を必要とせずに、幅広いSNNおよびTransformerベースのSNNモデルに即座に適用可能であると主張しています。

結論として、著者らは、不確実性ガイド型のサロゲート最適化が、特に低レイテンシ実行を必要とするアプリケーションにおいて、深層SNNの訓練を改善するための効果的な方向性であると述べています。