SNAP-UQ: TinyML 向けシングルパス・自己教師あり次活性化予測による不確実性推定

技術的サマリー（日本語）

本論文は、マイクロコントローラー（MCU）などのリソース制約の厳しいエッジデバイス向けに設計された、新しい不確実性推定手法SNAP-UQ（Self-supervised Next-Activation Prediction for single-pass Uncertainty）を提案するものです。TinyML モデルは、センサーのドリフトや分布シフト（OOD: Out-of-Distribution）に対して過剰に自信を持つ傾向があり、これによりデバイス上での信頼性の高いモニタリングや安全なフォールバックが困難になっています。既存の不確実性推定手法（深層アンサンブル、MC ドロップアウト、早期退出など）は、複数の推論パスや追加のメモリバッファ、大きな計算コストを必要とするため、ミリワット単位のハードウェアでは実用不可能な場合が多いです。SNAP-UQ は、これらの制約を克服し、単一のフォワードパスでラベルなしに不確実性を推定することを可能にします。

1. 課題背景

• TinyML の制約: 限られたフラッシュメモリ、RAM、計算能力、および厳格なレイテンシ制約。
• 既存手法の限界:
  • MC ドロップアウト/アンサンブル: 複数の推論パスが必要で、計算量とメモリ使用量が急増する。
  • 早期退出（Early Exit）: 追加のブランチが必要であり、softmax からの信号に依存するため、分布シフト下では脆い。
  • 従来の OOD 検出: 出力の自信度（softmax 確率）のみに依存しており、内部特徴量の分布シフトを捉えきれない場合がある。
• 目標: 単一パス、状態非依存（state-free）、ラベル不要、かつ MCU 上で実行可能な軽量な不確実性推定手法の確立。

2. 提案手法：SNAP-UQ

SNAP-UQ の核心は、ネットワークの**深度方向の次活性化予測（Next-Activation Prediction）**に基づいた「驚き（Surprisal）」信号の生成にあります。

2.1 基本的な仕組み

1. タップ（Tap）の選択: 既存のバックボーンネットワークの中間層（通常は中間ブロックと最終層直前のブロック）の少数の点（2〜3 箇所）を「タップ」として選択します。
2. 低ランク投影と予測ヘッド:
   • 各タップの直前の活性化 $a_{\ell-1}$ を、低ランクの射影（1x1 畳み込みや Skinny 線形層など）を用いて圧縮し、特徴量 $z_\ell$ を得ます。
   • 小さな予測ヘッド（Tiny Head）$g_\ell$ が、$z_\ell$ から次の層の活性化 $a_\ell$ の統計量（平均 $\mu_\ell$ と分散 $\sigma^2_\ell$）を予測します。
   • ヘッドは非常に小さく、INT8 量子化に対応しています。
3. 驚き（Surprisal）の計算:
   • 実際の活性化 $a_\ell$ と予測された分布との間の標準化された残差（Standardized Residual）を計算します。
   • 具体的には、$e_\ell = \|(a_\ell - \mu_\ell) \odot \sigma_\ell^{-1}\|^2_2$ となります。これは条件付き負対数尤度（Conditional NLL）に相当します。
   • この値が大きいほど、入力データが学習分布から外れている（ネットワークの内部ダイナミクスが予測と異なる）ことを示し、「驚き」が高いとみなされます。
4. 集約と較正:
   • 複数のタップで計算された驚きを重み付けして集約し、スカラー値 $S(x)$ を得ます。
   • この $S(x)$ を、オフラインで学習した単調なロジスティック関数またはアイソトニック回帰を用いて、0〜1 の範囲の不確実性スコア $U(x)$ に変換します。
   • 必要に応じて、分類器の即時の自信度（最大確率やマージン）と組み合わせて精度を向上させることも可能です。

2.2 技術的特徴と最適化

• 単一パス・状態非依存: 追加のフォワードパスや時間的バッファを必要とせず、通常の推論フローに組み込むだけで動作します。
• 整数演算対応（INT8）: ヘッドの重みは INT8 で量子化され、指数関数の計算はルックアップテーブル（LUT）で近似することで、MCU での実行を高速化・低消費電力化しています。
• 自己教師あり学習: ヘッドは、ラベルなしで「分布内（ID）データにおける層間遷移の予測可能性」を学習する補助タスクとして訓練されます。分類タスクとは独立して、ネットワークの内部ダイナミクスをモデル化します。
• アフィン不変性: バッチ正規化（BN）のようなチャネルごとの線形変換に対して、標準化された誤差は不変であることが理論的に保証されています。

3. 主要な貢献

1. 新しいパラダイム: 出力の自信度ではなく、「層から層への遷移の予測不可能性」に基づいた、TinyML 向けの単一パス不確実性推定手法を提案しました。
2. 極めて軽量な実装: 追加のフラッシュメモリは数十 KB 程度、レイテンシの増加は 2% 未満で済み、既存のアンサンブル手法や MC ドロップアウトに比べて、フラッシュサイズが 40〜60% 削減、レイテンシが 25〜35% 短縮されます。
3. 高い検出性能: 分布シフト（CID）下での精度低下検出、OOD 検出、および ID 内での誤分類検出において、既存の軽量手法を凌駕する性能を示しました。
4. 理論的裏付け: 提案手法が条件付きマハラノビス距離や負対数尤度と数学的に等価であることを示し、その不変性を証明しました。

4. 実験結果

MNIST、CIFAR-10、TinyImageNet、SpeechCommands などのデータセットと、2 種類の MCU（Big-MCU と Small-MCU）上で評価を行いました。

• リソース効率:
  • Small-MCU（フラッシュ < 1MB）: 既存のアンサンブル手法（Deep Ensemble, Early Exit Ensemble）はメモリ不足（OOM）で動作しませんでした。SNAP-UQ はすべてのタスクで動作し、Flash 使用量を 16〜24% 削減、レイテンシを 33% 短縮しました。
  • Big-MCU: 同様に、Flash 使用量を 40〜50% 削減、レイテンシを 24〜35% 短縮しました。
• 分布シフト検出（CID）:
  • 画像や音声のノイズ・歪み（MNIST-C, CIFAR-10-C, SpeechCmd-C）に対する精度低下の検出において、AUPRC（Average Precision-Recall Curve）が既存手法より高く、検出遅延が短縮されました。
  • 特に、シフトが進行するにつれて、SNAP-UQ のスコアが softmax 自信度よりも早期に反応することが確認されました。
• 失敗検出（OOD/ID 誤分類）:
  • ID 内での誤分類検出（ID✓— ID×）および OOD 検出（ID✓— OOD）において、AUROC が 0.9 前後と高い性能を維持しました。
  • 単一パスの決定論的手法としては、MC ドロップアウトや深層アンサンブルに匹敵する性能を達成しました。
• 較正性能:
  • 分布内データにおける期待較正誤差（ECE）や Brier スコアが改善され、より信頼性の高い確率推定が可能になりました。

5. 意義と結論

SNAP-UQ は、TinyML の実用化における重要な課題である「信頼性の高い不確実性推定」を、ハードウェア制約を大幅に超えることなく解決します。

• 実用性: 追加のラベルや複雑なアーキテクチャ変更を必要とせず、既存の TinyML モデルに容易に統合可能です。
• 安全性: 異常検知や安全なフォールバック（例：信頼性の低い推論を棄却し、人間への委譲や代替手段へ移行）を実現し、エッジ AI システムの安全性を向上させます。
• 将来展望: 本手法は、リソース制約の厳しい環境でも堅牢な AI モニタリングを実現する基盤技術として、産業応用（自動運転、医療、産業 IoT など）への展開が期待されます。

要約すると、SNAP-UQ は「ネットワークが自分の内部処理をどれだけ予測できるか」という自己教師的な視点を取り入れることで、リソース制約下でも高精度な異常検知を実現した画期的な手法です。