在標準的Transformer架構中，每個Token都會經過相同數量的層處理，無論其上下文難度如何。這種均勻處理方式導致計算資源浪費在簡單的Token上。針對這一問題，一篇來自arXiv的新論文提出了Token-Selective Attention（TSA），一種透過習得每Token門控來實現自適應計算深度的方法。
TSA的核心是在連續的兩個Transformer塊之間引入一個輕量級的門控機制。該門控是一個兩層多層感知器（MLP），它為一個殘差更新產生連續的停止機率。這使得整個機制端到端可微，且僅增加1.7%的引數，無需改變基礎架構。值得注意的是，即使在沒有任何深度正則化的情況下（λ=0），僅靠任務損失梯度驅動，TSA也能學習到跳過約20%的Token-層操作，表明其學會了與難度成比例的路由。
在字元級語言建模任務上，TSA在Tiny-Shakespeare和enwik8資料集上節省了14-23%的Token-層操作（TLOps），同時質量損失低於0.5%。在與早期退出方法匹配效率的條件下，TSA的驗證損失降低了0.7%。更重要的是，學習到的路由可以直接遷移到推理時的稀疏執行，從而實現實際的壁鐘時間加速效果。
TSA的門控機制不僅輕量，而且具有可解釋性。研究者發現，對於簡單的、資訊量低的Token，門控傾向於輸出較高的停止機率，從而跳過後續層的計算；而對於複雜的、資訊量高的Token，門控則允許其透過所有層。這種自適應行為完全由任務損失驅動，無需人為設定閾值。此外，TSA與現有的Transformer變體相容，可以輕鬆整合到預訓練模型中。實驗還表明，TSA的路由決策在訓練後可以固定下來，用於推理時的稀疏計算，從而顯著提升吞吐量。未來工作可以探索將TSA擴充套件到其他序列建模任務，如機器翻譯和文本生成，進一步驗證其通用性。該方法的提出者Ahmed Abdelmuniem Abdalla Mohammed在論文中詳細闡述了TSA的設計和實驗。該工作為Transformer的高效推理提供了新思路，尤其適合需要動態計算資源的場景。透過簡單的門控機制，TSA在幾乎不犧牲質量的前提下大幅降低了計算成本，有望在大型語言模型的部署中發揮重要作用。