上下文就緒的Transformer

近日，一篇題為“上下文就緒Transformer”的論文被提交至arXiv，並提出了一種創新的迴圈神經網路架構。該架構由D層Transformer模組和一個校正網路組成。與傳統Transformer不同，該架構在每個token進入模組之前，透過校正網路將之前位置的模組輸出（即快取的過去上下文摘要）與當前token嵌入結合，從而實現預上下文化。這使得token在進入模組時不再是原始嵌入，而是已經攜帶了上下文資訊。這種設計靈感來源於遞迴神經網路，但透過並行訓練解決了遞迴難以並行化的問題。

在訓練過程中，研究者將校正過程在完整序列上展開K次，每一步並行處理所有位置，使得訓練高效且可並行化。此外，任何預訓練的Transformer模型都可以透過新增一個零初始化的校正前饋網路並進行微調，轉換為上下文就緒模型。這意味著現有的預訓練模型可以輕鬆升級，而無需從頭訓練。

實驗表明，該架構在各種配置下均表現出色。在A100 GPU上，一個D=5的模型擊敗了12層的標準Transformer，生成速度提升1.7倍。更令人印象深刻的是，當K=10時，單層模型（D=1）以2.6倍的推理速度優勢超越了6層Transformer，且序貫推理的困惑度與並行K=10的版本僅相差0.01。這表明該架構不僅在效能上有所提升，還在推理效率上實現了顯著突破。

該架構特別受益於寬表示和長上下文。在指標追蹤任務中，使用BPTT訓練的D=1模型成功解決了所有10個組合層次，而標準Transformer則表現出階梯式的深度依賴，即需要更多層次才能處理更深的組合。這一突破錶明，透過預上下文化和校正機制，較淺的網路也能達到甚至超越深層網路的效能，為構建更高效的Transformer模型開闢了新方向。

相關論文已被NeurIPS接收，共22頁，作者為Mahesh Godavarti。