💻 computer science

Reg4Pru: Regularisation Through Random Token Routing for Token Pruning

本論文は、Vision Transformerにおけるトークン削減による性能低下を軽減する正則化手法であるReg4Pruを紹介しており、FIVES血管セグメンテーションデータセットにおいて、平均適合率の46%の絶対的な向上と29%の高速化を実現している。

原著者： Julian Wyatt, Ronald Clark, Irina Voiculescu

公開日 2026-02-04

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Julian Wyatt, Ronald Clark, Irina Voiculescu

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

Reg4Pru の解説：シンプルでクリエイティブな比喩を用いて

大きな問題：「混雑しすぎたパーティー」

Vision Transformer（現代の画像認識を支えるAIの脳）を、大規模で重要なパーティーだと想像してみてください。「ゲスト」はトークンです。これは、AIが分析している画像のごく小さな断片です。

標準的なセットアップでは、すべてのゲストが他のすべてのゲストと会話をして、画像が何であるかを判断します。これは精度面では素晴らしいのですが、信じられないほど時間がかかります。もしゲストの数が2倍になれば、必要な会話量は4倍になります。それは、スタジアムに詰めかけた人々の中で意味のある会話をしようとするようなもので、永遠に時間がかかってしまいます。

これを解決するために、研究者たちは**トークer Pruning（トークン・プルーニング／トークンの削減）**を考案しました。これは、パーティーをより速く進めるために、「退屈な」ゲスト（画像の冗長な部分）をドアの外へ追い出すドアマンのようなものです。

グリッチ（不具合）：「記憶喪失」効果

この論文は、このドアマン戦略における重大な問題を指摘しています。

学習中： AIはパーティー全体を見ながら学習し、時折、効率化のために人々を追い出す練習をします。
テスト時（推論時）： AIは時間を節定するために、ゲストを永久に追い出します。しかし、ここに落とし穴があります。詳細な決定（例えば、血管の詳細な地図を描くこと）を下さなければならないとき、AIは追い出したゲストを呼び戻して、欠けている詳細を補わなければなりません。

問題点： AIが混乱してしまうのです。「追い出された」ゲストたちは、残されたゲストたちがパーティーを楽しんで学習している間、暗闇の中に座っていました（レイヤーをスキップしていた）。そしてAIが、追い出されたゲストたちを再び光の中に連れ戻そうとしたとき、彼らはもう馴染めなくなっているのです。彼らは、自分が不在だった間に何が起きていたのかについての「記憶」を失っています（記憶喪失）。これにより、AIのパフォーマンスは、特にネットワークのより深く複雑なレイヤーにおいて、急激に低下します。それは、パズルのピースが途中で絵の様子を忘れてしまった状態で、複雑なパズルを完成させようとするようなものです。

解決策：Reg4Pru（「ランダムなリハーサル」）

著者たちは、Reg4Pru（Randomisation Through Random Token Routing／ランダムなトークン・ルーティングによる正則化）と呼ばれる新しい学習手法を提案しています。

これは、本番のショーの前にAIが行う**「リハーサル・ゲーム」**のようなものです。

単に人を追い出してうまくいくのを祈るのではなく、AIは学習中に、異なるグループのゲストに対して、ランダムな時間、パーティーの異なる部分を「スキップ」するようにランダムに指示するゲームを行います。

ランダムさが鍵： ある時はゲストAが2つの部屋をスキップし、別の時はゲストBが5つの部屋をスキップします。「スキップ・ゾーン」は毎回変化します。
結果： ゲストは常にランダムなセクションをスキップし、再び合流するように指示されるため、適応力を身につけます。彼らは、どこに参入したとしても、あるいはどれだけ長く離れていたとしても、再びパーティーに馴染み、最終的な絵に貢献できることを学びます。

この「ランダムなリハーサル」は、安定剤として機能します。これにより、ゲストを「暗闇」から連れ戻すことが安全であることをAIに教え、従来のメソッドで見られた混乱やパフォーマンスの低下を防ぎます。

結果：より速く、より賢く

チームは、目の画像から血管を見つけるという（高い詳細度を必要とするタスクである）FIVESという医療データセットを用いてテストを行いました。

比較： 彼らの手法を、標準的な「ドアマン」（この新しい学習なしのプルーニング）および「プルーニングなし」のベースラインと比較しました。
勝利：
- 精度： 標準的なプルーニング手法は、最終段階で非常にひどい結果（精度が大幅に低下）となりました。Reg4Pruはこれを修正し、標準的なプルーニング手法と比較して、平均精度を驚異的な**46%**向上させました。
- スピード： 彼らは、この高い精度を維持しながら、プルーニングなしの遅いバージョンよりも29%高速に処理することに成功しました。
- 視覚的効果： 論文ではヒートマップを用いて、標準的なプルーニング法がぼやけた乱れた結果を生むのに対し、Reg4Pruは、遅い非プルーニング版とほぼ同等の、鮮明でクリアな血管の画像を作成することを示しています。

まとめ

Reg4Pruは、画像の一部を無視することでスピードアップしようとする際に、AIが混乱するのを防ぐ学習テクニックです。学習中に「スキップ」と「再合流」をランダムに練習させることで、AIは安定性と正確さを維持することを学び、医師が必要とする微細な詳細を失うことなく、高解像度の医療画像をはるかに速く処理できるようになります。

技術要約: Reg4Pru – トークン・プルーニングのためのランダム・トークン・ルーティングによる正則化

問題提起
Vision Transformer (ViT) は、そのスケーラビリティと汎化能力により、現代のビジョンモデルの基礎となっています。しかし、その自己注意（self-attention）メカニズムへの依存は、トークン数 $N$ に対して二次関数的 ( $O(N^2)$ ) に増大する計算量という課題を抱えています。このボトルネックは、重要な診断詳細を保持するために高解像度の入力を必要とする医療画像セグメンテーションにおいて、特に深刻です。

トークン・プルーニング（冗長なトークンの削除）と再活性化戦略（高密度なピクセル単位の予測のために推論時に削除されたトークンを再導入すること）は効率化への道を示していますが、著者らは深い層における顕著な性能低下を観察しています。具体的には、推論時にプルーニングされたトークンを再活性化すると、それらが訓練時の状態と比較して劇的に異なる表現を示すことが挙げられます。この訓練時とテスト時の分布シフトは不安定性を引き起こし、モデルの深さが増すにつれて、直感に反して性能が低下するという現象をもたらします。

手法
本論文では、トークン・プルーニングと再活性化に伴う不安定性を軽減するために設計された、訓練時の正則化手法である Reg4Pru を提案しています。コアとなるメカニズムは、トークンを永久に破棄することなく、プルーニングの効果をシミュレートするために訓練フェーズで行われる ランダム・トークン・ルーティング です。

ランダム・トークン・ルーティング: 訓練時のフォワードパスにおいて、モデルは中間トランスフォーマー・ブロック間のランダムなルーティング範囲 $[l, n]$ を選択します。ランダムに選ばれたトークンのサブセット ( $\tau_r$ ) は、ブロック $l$ から $n$ までの計算をスキップし、残りのトークン ( $\tau_k$ ) は通常通り処理されます。
勾配の流れと表現: ルーティングされたトークンは特定のブロックをバイパスするため、それらの層に対する勾配が実質的に遮断されます。しかし、ステップごとにルーティングされるサブセットとルーティング範囲 $[l, n]$ の両方をランダム化することで、異なるトークンがイテレーションを通じて異なるブロックに勾配を寄与することになります。これにより、モデルは深さに依存しない表現を学習することを強制され、推論時に任意の深さでトークンを再活性化しても、表現のミスマッチに苦しむことがなくなります。
推論戦略: 推論時には、個別のポリシーネットワークが、階層的なキープ比率（keep-ratio）スケジュールに基づいてどのトークンをプルーニングするかを予測します。プルーニングされたトークンは、計算量を節約するために一時的に削除され、最終的なブロックで再活性化されて高密度なセグメンテーション・マスクを生成します。
損失関数: 訓練目的関数は、標準的な比率ベースの損失（トークン予算を強制するため）と、オプションとして、プルーニングの決定を導くために正解のセグメンテーション・マスクを使用するインフォームド・ポリシー損失 ( $L_{pol}$ ) を組み合わせたものです。ルーティング・メカニズムは、これらの損失とともに正則化器として機能します。

主な貢献

層ごとの不安定性の分析: 著者らは、高密度な予測におけるトークン・プルーニングにおける特定の不安定性（、すなわち、プルーニングされたトークンが訓練中にネットワークの全深度にさらされていないために、より深い層で性能が低下する現象）を特定し、特性を明らかにしました。
Reg4Pru フレームワーク: 彼らは、ランダムな境界を持つ訓練時のトークン・ルーティング戦略を導入しました。固定された境界や学習されたルーティング・ポリシーを使用する従来のルーティング手法とは異なり、Reg4Pru はランダムにルーティング範囲を変化させることで、あらゆるプルーニング深度における特徴分布を近似し、モデルを正則化します。
実証的検証: 本手法は、FIVES 網膜血管セグメンテーション・データセットを用いて検証されており、ランダム・ルーティングが固定ルーティングやルーティングなしと比較して、性能を大幅に安定させることを示しています。

結果
実験は、DINOv2-small で初期化された EoMT アーキテクチャを用い、FIVES データセット（2048 $\times$ 2048 画像をスループット・テスト用にクアドラント分割したもの）に対して行われました。

性能向上: ルーティングなしで訓練されたプルーニング・モデルと比較して、Reg4Pru は平均適合率（AP）を絶対値で 46% 向上させました（ $L_{ratio}$ 設定において 25.60% から 71.67% へ）。
効率性: Reg4Pru 設定は、高いセグメンテーション品質を維持しながら、非プルーニング・ベースラインと比較してウォークロック・タイムで 29% の相対的な高速化 を達成しました。
安定性: 層ごとの分析（図5）によれば、非ルーティングのプルーニング・モデルは深いブロックにかけて AP が急激に低下しましたが、Reg4Pru は安定した性能を維持しました。固定境界のルーティングはある程度の不安定性を緩和しましたが、固定範囲の終了後に発生する性能低下を防ぐことはできず、ランダム化の必要性を裏付けました。
比較: 同等のスループット構成において、特に ToMe が MHSA と MLP の両方の層に適用された場合、Reg4Pru は ToMe 手法よりも高い AP を示しました。

意義と主張
本論文は、Reg4Pru が高密度な予測タスクにおけるトークン削減戦略のための価値ある正則化器であることを主張しています。著者らは、プルーニングで見られる不安定性は、トークンの表現に関する訓練時とテスト時の分布シフトに起因すると考えています。訓練中にルーティング範囲をランダム化することで、Reg4Pru はモデルに対し、様々な深さでスキップされることに対して堅牢な表現を学習することを強制し、訓練と推論の間のギャップを効果的に埋めています。

著者らは、実験が高解像度の単一の医療データセットで行われたことを踏まえ、自らの知見の汎用性については慎重な姿勢を保っています。彼らは、本手法がセグメンテーションには有効である一方で、他の高密度な予測タスク（例：深度推定）への適用可能性や、ドロップアウトやストカスティック・デプスのような手法と組み合わせた際の一般的な正則化メカニズムとしての可能性を評価するには、さらなる検証が必要であると述べています。また、インフォームド・ポリシーについても検討しましたが、現在の設定では比率ベースのポリシーの方が優れた性能を示すことも認めています。