NVIDIA釋出Nemotron-Labs-TwoTower：基於凍結自迴歸骨幹網路的開放權重擴散語言模型

NVIDIA正式釋出了Nemotron-Labs-TwoTower，這是一種基於預訓練自迴歸骨幹網路的擴散語言模型。該模型以開放權重形式提供，採用NVIDIA Nemotron開放模型許可證。此次釋出旨在解決文本生成中的吞吐量瓶頸問題。

傳統的自迴歸（AR）模型逐token解碼，這種序列過程限制了生成吞吐量。而離散擴散語言模型則採用不同路徑：它們並行生成token並迭代最佳化。然而，大多數擴散語言模型使用單一網路同時處理兩項任務：在每個步驟中表示乾淨token並去噪被汙染的token。TwoTower將這兩項任務分離到兩個獨立模組中，分別稱為“上下文塔”和“去噪塔”。

TwoTower基於Nemotron-3-Nano-30B-A3B骨幹網路構建，該骨幹網路是開放權重的混合架構，融合了Mamba-2、自注意力和混合專家（MoE）層。每個塔包含52層：23個Mamba-2層、6個自注意力層和23個MoE層。釋出的檢查點包含兩個塔，總計約600億引數，每個塔每token有效引數約30億。MoE使用128個可路由專家，其中6個啟用，外加2個共享專家。

兩個塔初始化為同一骨幹檢查點的副本，但僅去噪塔進行訓練，AR上下文塔保持凍結。去噪器的訓練資料量約為2.1萬億token，僅是骨幹網路預訓練資料（25萬億token）的一小部分。

在推理過程中，AR上下文塔對提示和已提交token進行因果處理，生成逐層的KV快取和最終的Mamba-2狀態，保留了骨幹網路的自迴歸能力。擴散去噪塔則對噪聲塊進行迭代最佳化：在塊內使用雙向注意力，同時對過去的乾淨塊保持因果性。兩個塔逐層連線，去噪塔的每一層都交叉關注上下文塔的對應層，這種逐層對齊的交叉注意力提供了對骨幹網路表示的多尺度訪問。此外，Mamba-2層從上下文塔的Mamba狀態初始化其初始狀態，擴散時間步透過adaLN-single時間調節調變每一層，該模組僅增加約150萬引數。

生成過程按塊進行：每個塊初始化為S個[MASK]token，去噪器經過T步迭代後提交該塊，然後上下文塔處理已提交的token以更新其快取。這解釋了為什麼多次去噪步驟仍能快於逐token解碼：自迴歸解碼每步只提交一個token，而TwoTower在早期迭代中每步提交多個token。

在2×H100 GPU上的BF16評估顯示，預設操作點（置信度解掩碼，閾值γ=0.8，塊大小S=16）下，TwoTower在MMLU、ARC-Challenge、WinoGrande等任務上效能接近AR基線，聚合質量保留98.7%，生成吞吐量提升2.42倍。通用知識任務與基線差距在1個百分點以內，程式碼和數學任務下降稍多，常識和多語言任務則略有恢復或提升。降低γ可提交更多token以提升吞吐量，但會帶來質量損失。

該檢查點支援三種推理模式：完整雙塔擴散（需2塊GPU，每塊約59GB BF16記憶體）、模擬自迴歸和純自迴歸（後者可在單塊80GB GPU上執行）。最直接的應用場景是加速批次生成，例如合成文本生成任務中，可以以少量質量下降換取大幅吞吐量提升。使用者還可調節置信度閾值來權衡質量與速度。此外，由於上下文塔保留了LM head，可用於推測解碼、驗證或自迴歸評分，團隊可從同一檢查點執行多種模式。

TwoTower的主要優勢包括：開放權重，可用於商業用途；在預設操作點下以2.42倍吞吐量保留98.7%質量；單個檢查點支援多種解碼模式；去噪器訓練成本較低。其侷限包括：完整雙塔推理需要2塊GPU；程式碼和數學任務效能下降較多；固定模型權重記憶體佔用較高；釋出版本為基礎模型，未經過指令微調或對齊。