超越Transformer與Mamba，Meta聯合斯坦福等高校推出最強架構TTT

作者 | 趙明華

近日，斯坦福、UCSD、UC 伯克利和 Meta 的研究人員提出了一種全新架構，用機器學習模型取代 RNN 的隱藏狀態。

圖 1 所有序列建模層都可以表示為一個根據更新規則轉換的隱藏狀態

這個模型透過對輸入 token 進行梯度下降來壓縮上下文，這種方法被稱為「測試時間訓練層（Test-Time-Training layers，TTT）」。該研究作者之一 Karan Dalal 表示，他相信這將根本性地改變語言模型方法。

自注意力機制在處理長上下文時表現良好，但其複雜度是二次的。現有的 RNN 層具有線性複雜度，但其在長上下文中的表現受限於其隱藏狀態的表達能力。隨著上下文長度的增加，成本也會越來越高。

作者提出了一種具有線性複雜度和表達能力強的隱藏狀態的新型序列建模層。關鍵思路是讓隱藏狀態本身成為一個機器學習模型，並將更新規則設為自監督學習的一步。

圖 2，RNN 層與 TTT 層的關係

論文中提出了兩種例項：TTT-Linear 和 TTT-MLP，它們的隱藏狀態分別是線性模型和兩層 MLP。團隊在 125M 到 1.3B 引數規模上評估了例項，並與強大的 Transformer 和現代 RNN Mamba 進行了比較。結果顯示，與 Mamba 相比，TTT-Linear 的困惑度更低，FLOP 更少（左），對長上下文的利用更好（右）：

圖 3 TTT-Linear 與 Mamba 對比

這個結果代表了現有 RNN 的尷尬現實。一方面，RNN（與 Transformer 相比）的主要優點是其線性（與二次型）複雜性。這種漸近優勢只有在長上下文的實踐中才能實現，根據下圖，這個長度是 8k。另一方面，一旦上下文足夠長，現有的 RNN（如 Mamba）就很難真正利用所依賴的額外資訊。

圖 4 TT Linear 在 8k 環境下已經比 Transformer 更快

並且，大量的實驗結果表明：TTT-Linear 和 TTT-MLP 都匹配或超過基線。與 Transformer 類似，它們可以透過限制更多的代幣來不斷減少困惑，而 Mamba 在 16k 上下文後則不能。經過初步的系統最佳化，TTT Linear 在 8k 環境下已經比 Transformer 更快，並且在 wall-clock 時間上與 Mamba 相匹配。

TTT 層在理論上和實驗評估中表現出色，尤其是在長上下文處理和硬體效率方面。如果在實際應用中能夠解決一些潛在的工程挑戰，如大規模部署和整合問題，工業界對 TTT 層的接受度也將逐步提升。

論文連結：https://arxiv.org/pdf/2407.04620v1

內容推薦

AIGC技術正以驚人的速度重塑著創新的邊界，InfoQ 首期《大模型領航者AIGC實踐案例集錦》電子書，深度對話30位國內頂尖大模型專家，洞悉大模型技術前沿與未來趨勢，精選10餘個行業一線實踐案例，全面展示大模型在多個垂直行業的應用成果，同時，揭秘全球熱門大模型效果，為創業者、開發者提供決策支援和選型參考。關注「AI前線」，回覆「領航者」免費獲取電子書。

活動推薦

AICon 全球人工智慧開發與應用大會，為資深工程師、產品經理、資料分析師等專業人群搭建深度交流平臺。聚焦大模型訓練與推理、AI Agent、RAG 技術、多模態等前沿議題，匯聚 AI 和大模型超全落地場景與最佳實踐，期望幫助與會者在大模型時代把握先機，實現技術與業務的雙重飛躍。

在主題演講環節，我們已經邀請到了「蔚來創始人李斌」，分享基於蔚來汽車 10 年來創新創業過程中的思考和實踐，聚焦 SmartEV 和 AI 結合的關鍵問題和解決之道。大會火熱報名中，7 月 31 日前可以享受 9 折優惠，單張門票節省 480 元（原價 4800 元），詳情可聯絡票務經理 13269078023 諮詢。