NVIDIA「雙塔」擴散模型問世：Nemotron-Labs-TwoTower 如何在不犧牲品質下，讓文字生成速度飆升 2.42 倍？

人工智慧領域巨擘 NVIDIA 近期發布了 Nemotron-Labs-TwoTower，這是一款基於預訓練自迴歸（Autoregressive, AR）骨幹模型所打造的擴散式語言模型。這項創新以開源權重形式釋出，採用 NVIDIA Nemotron 開源模型授權，其核心目標是突破現有文字生成模型在「吞吐量」上的瓶頸。

突破自迴歸模型的速度限制：擴散式架構的優勢

傳統的自迴歸模型在生成文字時，是逐一解碼每個詞元（token），這種序列化的處理方式，限制了文字生成的吞吐量。相較之下，離散擴散式語言模型則採取不同的路徑，它們能以平行方式生成詞元，並透過迭代精煉來逐步完善輸出。

然而，大多數擴散式語言模型通常只使用一個網路來執行兩項任務：既要表示「乾淨」的詞元，又要在每個步驟中對「受損」的詞元進行去噪。Nemotron-Labs-TwoTower 的關鍵創新點，便在於將這兩項任務分離，劃分為兩個獨立的「塔」（tower）來執行。這種「雙塔」設計，不僅保留了自迴歸基準模型高達 98.7% 的綜合基準品質，更將實際的文字生成吞吐量提升了 2.42 倍。

Nemotron-Labs-TwoTower 的核心效益：

• 品質保留：98.7% 的自迴歸模型基準品質。
• 生成吞吐量：提升 2.42 倍（在 γ=0.8、S=16、2 張 H100 GPU 的條件下）。

「雙塔」如何運作？技術解密

Nemotron-Labs-TwoTower 是一個區塊式（block-wise）自迴歸擴散模型，其基礎建構在 Nemotron-3-Nano-30B-A3B 這個開源權重的混合骨幹模型之上。這個骨幹模型巧妙地結合了 Mamba-2、自注意力機制（self-attention）以及專家混合層（Mixture-of-Experts, MoE）。

每個「塔」都包含 52 層，其中有 23 層 Mamba-2、6 層自注意力機制，以及 23 層 MoE。NVIDIA 釋出的檢查點（checkpoint）包含了這兩個塔，總參數約為 600 億。每個詞元在每個塔中啟用的參數約為 30 億。MoE 層則使用了 128 個可路由的專家，其中 6 個會被啟用，並額外包含 2 個共享專家。

值得注意的是，兩個塔最初都是同一個骨幹模型檢查點的副本，但只有「去噪器塔」（denoiser tower）會進行訓練，而「自迴歸上下文塔」（AR context tower）則保持凍結狀態。去噪器塔約使用了 2.1 兆個詞元進行訓練，這僅是骨幹模型預訓練時所用 25 兆個詞元的一小部分。

雙塔協作機制

• 自迴歸上下文塔：負責對提示（prompt）和已提交的詞元進行因果（causal）處理，產生每層的鍵值快取（KV cache）和最終的 Mamba-2 狀態，並保留骨幹模型的自迴歸能力。
• 擴散去噪器塔：負責精煉嘈雜的區塊。在每個區塊內部，它使用雙向的區塊內注意力機制，並對過去的「乾淨」區塊保持因果關係。

這兩個塔是逐層連接的。去噪器塔的第 i 層會與上下文塔的第 i 層進行交叉注意力（cross-attend），這種層級對齊的交叉注意力機制，讓模型能夠多尺度地存取骨幹模型的表示。相較於以往僅廣播最後一個隱藏狀態的方法，這是一個顯著的改進。

此外，去噪器塔還有兩項重要的修改：Mamba-2 層的初始狀態會從上下文塔的 Mamba 狀態中獲取，而擴散時間步長則透過 AdaLN 單一時間條件化（adaLN-single time conditioning）模組來調節每一層，這個 AdaLN 模組僅增加了約 150 萬個參數。

文字生成過程是區塊式進行的。每個區塊最初由 S 個 [MASK] 詞元組成，去噪器塔會在 T 個步驟中對其進行精煉，然後提交。隨後，上下文塔會處理這些已提交的詞元，以更新其快取。這也解釋了為何多個去噪步驟仍能超越單詞元解碼的速度，因為自迴歸解碼每一步只提交一個詞元，而 TwoTower 在精煉初期每一步就能提交多個詞元。

效能基準與彈性操作模式

在 BF16 浮點數格式下，使用兩張 H100 GPU 進行的評估顯示，Nemotron-Labs-TwoTower 在預設操作點（信心去遮蔽，閾值 γ=0.8，區塊大小 S=16）的表現如下：

在通用知識任務上，其品質與自迴歸基準模型幾乎保持一致。程式碼和數學任務則顯示出輕微的品質下降，而常識和多語言任務的得分則有所恢復或略有提升。值得注意的是，降低 γ 值可以讓模型在每一步提交更多詞元，進而提高吞吐量，但同時也會導致品質略微下降。

三種生成模式

Nemotron-Labs-TwoTower 的檢查點提供了三種推論路徑：

1. 全雙塔擴散模式（generate_mask_diffusion()）：需要兩張 GPU，每張 GPU 約需 59GB 的 BF16 記憶體，每一步最多可提交 block_size 個詞元。
2. 模擬自迴歸模式（generate_mock_ar()）：每一步提交一個詞元。
3. 純自迴歸模式（generate_ar()）：可在單張 80GB GPU 上運行，每一步提交一個詞元。

應用場景與未來展望

Nemotron-Labs-TwoTower 最直接的應用場景是加速批次生成。對於需要大量生成合成文本的資料團隊而言，可以透過犧牲微小的品質（在 γ=0.8 時，約 1.3% 的品質損失）來換取 2.42 倍的生成速度提升。

其次，它提供了調整品質與吞吐量之間權衡的彈性。提高 γ 值可以保留更多品質，而降低 γ 值則能讓每一步提交更多詞元，以追求更高的速度。

第三個應用是「即插即用」的適應性。上下文塔保留了其語言模型頭部，可用於推測性解碼（speculative decoding）、驗證或自迴歸評分。這意味著團隊可以從同一個檢查點同時運行自迴歸和擴散模式。

優勢與劣勢

優勢：

• 在 NVIDIA Nemotron 開源模型授權下，開源權重可供商業使用。
• 在預設操作點下，以 2.42 倍的吞吐量保留了 98.7% 的自迴歸品質。
• 一個檢查點支援擴散、模擬自迴歸和純自迴歸三種解碼模式。
• 去噪器僅訓練了約 2.1 兆個詞元，而非完整的重新預訓練。
• 序列長度快取記憶體擴展方式與自迴歸基準模型相同。

劣勢：

• 完整的雙塔擴散模式需要兩張 GPU，每張約 59GB 的 BF16 記憶體。
• 程式碼和數學任務的品質下降幅度略大於通用知識（例如 HumanEval 從 79.27 下降到 75.58）。
• 同時駐留兩個塔會增加固定的模型權重記憶體佔用。
• 釋出的檢查點是一個基礎模型，尚未經過指令微調（instruction tuning）或對齊（alignment）。
• 吞吐量超過 3 倍時，會伴隨更大的品質損失。

總體而言，NVIDIA Nemotron-Labs-TwoTower 的推出，為大型語言模型的部署和研究帶來了新的可能性。它透過創新的雙塔架構，在效率與品質之間找到了新的平衡點，特別是在需要快速生成大量文本的應用場景中，展現了其獨特的價值。這項技術的開源，也將進一步加速相關領域的創新與應用。