人工智能芯片巨頭英偉達的研究人員近日發佈了一項名爲“FFN融合”（FFN Fusion）的創新架構優化技術。該技術旨在通過解決Transformer架構中固有的串行計算瓶頸，顯著提升大型語言模型（LLMs）的推理效率，爲更廣泛地部署高性能AI應用鋪平道路.

近年來，大型語言模型在自然語言處理、科學研究和對話代理等領域展現出強大的能力。然而，隨着模型規模和複雜性的不斷增加，其推理過程所需的計算資源也大幅增長，導致了效率瓶頸。Transformer架構是LLM的基礎，其交替的注意力機制和前饋網絡（FFNs）層需要按順序處理輸入。這種固有的串行結構在模型規模擴大時，會顯著增加計算和GPU之間的通信成本，降低效率並提高部署成本。尤其是在需要快速生成多個token的場景（如實時AI助手）中，這個問題更加突出。

QQ_1743400012665.png

爲了應對這一挑戰，英偉達的研究人員提出了FFN融合技術。該方法的核心思想是將模型中連續的、相互依賴性較低的FFN層合併爲一個更寬的FFN。研究人員觀察到，在移除注意力層後，LLM中通常會存在較長的連續FFN序列。通過分析這些序列，他們發現這些FFN層之間的依賴性很小，因此可以並行執行。

FFN融合的數學基礎在於將多個串聯FFN的權重進行拼接，從而創建一個等效的、可以並行計算的單一模塊。例如，如果三個FFN順序堆疊，每個FFN的輸出作爲下一個的輸入，FFN融合則消除了這種依賴關係，使得這三個FFN能夠同時處理相同的輸入，並將其輸出進行聚合。理論分析表明，融合後的FFN保持了與原始FFN相同的表示能力。

Ultra-253B-Base:性能與效率的雙重提升

英偉達的研究人員將FFN融合技術應用於Meta的Llama-3.1-405B-Instruct模型，通過剪枝和重構創建了一個名爲Ultra-253B-Base的新模型。實驗結果表明，Ultra-253B-Base在推理速度和資源效率方面取得了顯著的提升。具體來說，該模型在批量大小爲32時，推理延遲降低了1.71倍，每個token的計算成本降低了35倍。

更令人印象深刻的是，效率的提升並沒有以犧牲模型能力爲代價。Ultra-253B-Base在多個權威評測基準上取得了優異的成績，例如:MMLU85.17%，MMLU-Pro72.25%，HumanEval86.58%，Arena Hard84.92%，MT-Bench9.19。這些結果通常與原始的4050億參數模型相當甚至更高，而Ultra-253B-Base僅包含2530億參數。此外，該模型的內存使用量也減少了一半，這得益於kv-cache的優化。

QQ_1743400044821.png

研究人員使用餘弦距離分析FFN層之間的輸出，以識別低相互依賴性的區域，這些區域是進行融合的最佳選擇。FFN融合技術已在不同規模的模型（包括490億、700億和2530億參數）上得到驗證，表明其具有良好的通用性。

這項研究表明，通過深入分析和巧妙的架構設計，可以顯著提升LLM的效率。FFN融合爲設計更並行化、更適應硬件的LLM奠定了基礎。雖然完全Transformer模塊的並行化由於更強的層間依賴性而面臨更多挑戰，但FFN融合的成功無疑爲未來LLM的效率優化指明瞭一個重要的方向。

論文:https://arxiv.org/abs/2503.18908