標準的なTransformerアーキテクチャでは、各トークンはコンテキストの難易度に関わらず同じ数の層を通過します。この均一な処理は、簡単なトークンに対して計算リソースを無駄に消費します。この問題を解決するために、arXivに投稿された新しい論文では、Token-Selective Attention（TSA）という手法を提案しています。TSAは、学習可能なトークン単位のゲートを残差更新に導入し、連続するTransformerブロック間で適応的に計算深度を制御します。
各ゲートは軽量な2層の多層パーセプトロン（MLP）で構成され、連続的な停止確率を出力します。これにより、機構全体がエンドツーエンドで微分可能となり、ベースアーキテクチャを変更せずにパラメータの増加はわずか1.7%に抑えられます。注目すべき点は、深さの正則化がない場合（λ=0）でも、タスク損失の勾配だけでルーターが約20%のトークン-層操作をスキップするように学習されることです。これは、ルーターが難易度に比例したルーティングを自動的に獲得することを示しています。
文字レベル言語モデリングの実験では、TSAはTiny-Shakespeareとenwik8データセットで14-23%のトークン-層操作（TLOps）を削減し、品質損失は0.5%未満でした。同じ効率で比較した場合、TSAは早期終了法よりも検証損失を0.7%低減しました。さらに、学習されたルーティングは推論時のスパース実行に直接転用でき、実際の壁時計時間での高速化を実現します。
TSAのゲートメカニズムは軽量であるだけでなく、解釈可能性も備えています。研究者は、単純で情報量の少ないトークンに対してはゲートが高い停止確率を出力してその後の層をスキップし、複雑で情報量の多いトークンに対しては全ての層を通過させることを発見しました。この適応的な動作はタスク損失のみによって駆動され、人為的な閾値設定を必要としません。また、TSAは既存のTransformer変種と互換性があり、事前学習済みモデルに容易に統合できます。実験では、TSAのルーティング決定は学習後に固定でき、推論時のスパース計算に利用することでスループットを大幅に向上させることが示されました。将来の研究では、TSAを機械翻訳やテキスト生成などの他の系列モデリングタスクに拡張し、その汎用性を検証することが期待されます。本手法の提案者であるAhmed Abdelmuniem Abdalla Mohammed氏は、論文でTSAの設計と実験を詳細に説明しています。この研究は、Transformerの効率的な推論に新たな道を開き、特に動的な計算リソースが求められるシナリオで有用です。単純なゲート機構により、TSAは品質をほとんど犠牲にすることなく計算コストを大幅に削減し、大規模言語モデルの展開において重要な役割を果たす可能性があります。