上下文就緒的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。