DANTINOX：用于多范式语言建模的统一框架

# 面向多范式语言建模的统一框架
来源：https://arxiv.org/html/2609.17535
Marco Simoni, Aleksandar Fontana, Giulio Rossolini, Andrea Saracino — TeCIP机器人与人工智能卓越部门，比萨圣安娜高等研究院

###### 摘要
语言生成研究日益涵盖三大范式：自回归解码、离散掩码扩散和连续流匹配。由于每种范式都存在于独立的代码库中，因此比较它们十分困难，测量出的差异通常反映的是实现细节而非范式本身。我们推出**DantinoX**，一个开源的JAX/Flax库，其中单一的模块化Transformer骨干架构服务于所有三种范式。切换生成范式、注意力机制或硬件拓扑仅需更改配置，而骨干架构、分词器、初始化策略和训练基础设施保持一致。这使得在一个统一的API内进行训练、流式推理和基准测试的受控跨范式比较成为可能。利用其自身的流水线，我们评估了生成质量和注意力组合，并通过硬件屋顶线分析剖析了推理延迟、吞吐量和能耗。**DantinoX**根据MIT许可证发布，可通过`pip`安装，代码、文档、notebooks和演示视频均已公开<sup>111</sup>。

**安装包**（https://pypi.org/project/dantinox/）— **演示视频**（https://youtu.be/1u5-AieDzIc）

## 1 引言
语言建模不再等同于从左到右的解码。除了占主导地位的自回归（AR）范式 [Brown et al., 2020 (https://arxiv.org/html/2609.17535#bib.bib2)] 之外，掩码扩散 [Austin et al., 2021 (https://arxiv.org/html/2609.17535#bib.bib3); Li et al., 2022 (https://arxiv.org/html/2609.17535#bib.bib4)] 和连续流匹配 [Hu et al., 2026 (https://arxiv.org/html/2609.17535#bib.bib5)] 也是可行的替代方案 [Zhao et al., 2026 (https://arxiv.org/html/2609.17535#bib.bib1)]，它们在训练目标、注意力掩码、解码过程和推理成本上各有不同。然而，对它们进行实证比较具有挑战性 [Dodge et al., 2019 (https://arxiv.org/html/2609.17535#bib.bib77)]。每种范式都发展了自己的软件生态系统，而生产框架仍然主要为AR推理而优化 [Shoeybi et al., 2019 (https://arxiv.org/html/2609.17535#bib.bib6)]。研究者若要比较扩散模型与AR基线，必须将模型移植到不兼容的仓库中，而他们测量出的任何差异可能反映的是分词器、初始化或训练循环细节，而非范式本身。所缺失的是一个共享的基础平台，使得不同范式能在相同条件下进行训练、比较和服务 [Dodge et al., 2019 (https://arxiv.org/html/2609.17535#bib.bib77)]。

我们推出 **DantinoX**，一个提供此平台的开源库。它基于JAX/Flax构建 [Bradbury et al., 2018 (https://arxiv.org/html/2609.17535#bib.bib7); Heek et al., 2024 (https://arxiv.org/html/2609.17535#bib.bib8)]，充分利用了JAX的函数式编程模型和可组合转换。其模块化的Transformer骨干架构通过简单配置支持所有三种范式，为整个训练、评估和服务生命周期提供统一的接口。该骨干在注意力机制、前馈层、LoRA适配器 [Hu et al., 2022 (https://arxiv.org/html/2609.17535#bib.bib42)] 和数据/张量并行 [Adler et al., 2024 (https://arxiv.org/html/2609.17535#bib.bib61)] 方面具有高度的可组合性，并包含一套集成的基准测试工具。

**DantinoX** 旨在服务于广泛的用户群体：研究人员，可通过配置更改而非模型代码编辑实现消融研究；教育者，提供了一个易于访问且统一的代码库，用于教授多范式自然语言处理；以及小型企业的从业者，提供了可复用的流水线，减少了从头构建此类基础设施的需求 [Weeger et al., 2025 (https://arxiv.org/html/2609.17535#bib.bib13); Hansen and Bøgh, 2021 (https://arxiv.org/html/2609.17535#bib.bib14); Schwaeke et al., 2025 (https://arxiv.org/html/2609.17535#bib.bib15)]。该库采用MIT许可证，代码、文档和配置可在GitHub上获取<sup>222</sup>。我们通过三个评估套件（第4节）对其进行验证：与其它代码库 [Zhou et al., 2026 (https://arxiv.org/html/2609.17535#bib.bib72); Patel et al., 2026 (https://arxiv.org/html/2609.17535#bib.bib71)] 进行范式实现的交叉检查；在所有范式-注意力组合上评估生成质量（MAUVE [Pillutla et al., 2021 (https://arxiv.org/html/2609.17535#bib.bib64)]、困惑度、多样性、条件BLEU [Papineni et al., 2002 (https://arxiv.org/html/2609.17535#bib.bib65)]）；以及评估推理效率。

**表1：DantinoX与现有NLP生成建模框架的比较。**
| 框架         | 生态系统         | 范式支持    |              |              | 注意力变体                 | LLM基础设施 |              |              |
|--------------|------------------|-------------|--------------|--------------|----------------------------|-------------|--------------|--------------|
|              |                  | AR          | 离散         | 连续         |                            | LoRA        | 多GPU        | 基准套件     |
| HuggingFace  | PyTorch / JAX    | ✓           | ✗            | ✗            | MHA, GQA, MLA               | ✓           | ✓            | ✗            |
| MaxText      | JAX / Flax       | ✓           | ✗            | ✗            | MHA, GQA                   | ✓           | ✓            | ✗            |
| Levanter     | JAX / Flax       | ✓           | ✗            | ✗            | MHA, GQA                   | ✓           | ✓            | ✗            |
| OpenLM       | PyTorch          | ✓           | ✗            | ✗            | MHA                        | ✗           | ✓            | ✗            |
| torchtune    | PyTorch          | ✓           | ✗            | ✗            | MHA, GQA                   | ✓           | ✓            | ✗            |
| Fairseq      | PyTorch          | ✓           | ✗            | ✗            | MHA                        | ✗           | ✓            | ✗            |
| xLM          | PyTorch          | ✓           | ✓            | ✗            | MHA                        | ✗           | ✓            | ✗            |
| dLLM         | PyTorch          | ✗           | ✓            | ✗            | MHA                        | ✓           | ✓            | ✗            |
| **DantinoX** | **JAX / Flax**   | **✓**       | **✓**        | **✓**        | **MHA, GQA, MLA**          | **✓**       | **✓**        | **✓**        |

## 2 相关工作
语言建模框架沿着AR/非AR的边界划分为两派，以灵活性换取规模。面向生产的库，如Hugging Face `transformers` [Wolf et al., 2020 (https://arxiv.org/html/2609.17535#bib.bib20)]、Megatron-LM [Shoeybi et al., 2019 (https://arxiv.org/html/2609.17535#bib.bib6)]、Nanotron [Face, 2023 (https://arxiv.org/html/2609.17535#bib.bib69)] 和 MaxText [Google Cloud, 2023 (https://arxiv.org/html/2609.17535#bib.bib22)] 提供了高度优化的大规模训练基础设施，但它们的生成和解码流水线假设的是从左到右的因果推理。第二波以研究为中心的框架则倾向于模块性和可复现性，而非原始规模：Levanter [Haller et al., 2023 (https://arxiv.org/html/2609.17535#bib.bib67)] 和 OpenLM [Gururangan et al., 2023 (https://arxiv.org/html/2609.17535#bib.bib70)] 针对可复现的AR预训练，而 torchtune [Team, 2024 (https://arxiv.org/html/2609.17535#bib.bib66)] 专注于AR后训练和微调。这三者仍然局限于AR范式，需要对其核心API和注意力机制进行侵入式更改，才能支持掩码扩散或连续流匹配等非自回归目标。Fairseq [Ott et al., 2019 (https://arxiv.org/html/2609.17535#bib.bib23)] 通过非自回归翻译提供了对并行解码的早期支持，但从未将其扩展到连续空间目标或现代注意力变体。非自回归库仅部分地弥补了这一空白。xLM [Patel et al., 2026 (https://arxiv.org/html/2609.17535#bib.bib71)] 在共享骨干上支持AR基线以及离散掩码扩散，但不支持连续流匹配。dLLM [Zhou et al., 2026 (https://arxiv.org/html/2609.17535#bib.bib72)] 针对诸如LLaDA [Nie et al., 2026 (https://arxiv.org/html/2609.17535#bib.bib25)] 和DreamYeet [Ye et al., 2025a (https://arxiv.org/html/2609.17535#bib.bib73)] 等模型的离散扩散变体（掩码、基于块、基于编辑），但没有在同一骨干上提供匹配的AR基线。简而言之，没有现有框架能在单一骨干上统一AR、离散扩散和连续流匹配的训练，迫使跨范式研究调和那些掩盖了真实算法差异的异构实现。

**DantinoX** 旨在弥合这一差距。如表1所示，它在单一的JAX/Flax骨干上结合了所有三种范式，同时还包括了注意力变体（MHA、GQA、MLA）、LoRA微调、多GPU扩展和基准测试基础设施——而这些现有工具往往只是孤立地提供。切换范式、注意力机制或硬件拓扑仅需更改配置，而非一个新代码库，让研究人员和从业者能在同一API内工作。

## 3 DantinoX 架构
评估文本生成范式（如自回归、离散扩散和流匹配）需要共享的架构和训练设置。当方法使用不同代码库时，实现细节可能会隐藏算法行为。**DantinoX** 通过将模型骨干与生成方法分离来解决此问题。核心模型，或称**骨干**，由可互换组件（如注意力、前馈层和归一化）构建而成。用户可以通过一个配置设置来选择生成范式。

### 3.1 骨干
为确保跨范式的受控比较，共享骨干（图1 [https://arxiv.org/html/2609.17535#S3.F1]）被实现为一个高度模块化的预归一化Transformer [Radford et al., 2019 (https://arxiv.org/html/2609.17535#bib.bib26); Xiong et al., 2020 (https://arxiv.org/html/2609.17535#bib.bib27)]，包含`L`个可配置的块。**DantinoX** 并非强制静态架构，而是将关键结构选择作为超参数开关暴露出来。

对于注意力子层，该框架支持标准的MHA [Vaswani et al., 2017 (https://arxiv.org/html/2609.17535#bib.bib9)] 作为参考基线，以及诸如GQA [Ainslie et al., 2023 (https://arxiv.org/html/2609.17535#bib.bib29)] 和MLA [Liu et al., 2024 (https://arxiv.org/html/2609.17535#bib.bib30)] 等内存高效变体。所有注意力机制（根据**范式**，可以是**因果的**或**双向的**）都无缝集成了现代优化，包括FlashAttention [Dao et al., 2022 (https://arxiv.org/html/2609.17535#bib.bib31)]、滑动窗口上下文 [Beltagy et al., 2020 (https://arxiv.org/html/2609.17535#bib.bib32)]、门控 [Qiu et al., 2026 (https://arxiv.org/html/2609.17535#bib.bib75)]、线性 [Katharopoulos et al., 2020 (https://arxiv.org/html/2609.17535#bib.bib80)] 和差分注意力 [Ye et al., 2025b (https://arxiv.org/html/2609.17535#bib.bib33)]。同样，前馈网络（FFN）可以实例化为带有SwiGLU [Shazeer, 2020 (https://arxiv.org/html/2609.17535#bib.bib34)] 或其他激活函数（如GELU [Hendrycks and Gimpel, 2016 (https://arxiv.org/html/2609.17535#bib.bib35)]）的标准密集MLP、带有top-k路由的稀疏混合专家（MoE）[Shazeer et al., 2017 (https://arxiv.org/html/2609.17535#bib.bib36); Fedus et al., 2022 (https://arxiv.org/html/2609.17535#bib.bib37)]，或最近的潜在MoE [Elango et al., 2026 (https://arxiv.org/html/2609.17535#bib.bib49)]。用户可以选择位置编码方案（RoPE [Su et al., 2024 (https://arxiv.org/html/2609.17535#bib.bib38)]、正弦 [Vaswani et al., 2017 (https://arxiv.org/html/2609.17535#bib.bib9)]、学习式或无）以及归一化层（RMSNorm [Zhang and Sennrich, 2019 (https://arxiv.org/html/2609.17535#bib.bib40)] 或 LayerNorm [Ba et al., 2016 (https://arxiv.org/html/2609.17535#bib.bib41)]）。为了进行下游适配，可以将LoRA适配器 [Hu et al., 2022 (https://arxiv.org/html/2609.17535#bib.bib42)] 注入任何投影矩阵。此外，对于自回归和离散扩散建模，原生支持输入嵌入和输出投影之间的权重共享 [Press and Wolf, 2017 (https://arxiv.org/html/2609.17535#bib.bib47); Inan et al., 2016 (https://arxiv.org/html/2609.17535#bib.bib48)]。

**图1：统一骨干。**  
输入：Token ID  $\mathbf{x} \in \mathbb{Z}^{B \times T}$ → 嵌入层 + 位置编码：RoPE | 正弦 | 学习式 → 预归一化（RMSNorm | LayerNorm） → 自注意力：MHA | GQA | MLA （可附加：Flash · 滑动窗口 · 差分 · 门控 · 线性） → 预归一化（RMSNorm | LayerNorm） → 前馈网络：密集（MLP, SwiGLU） | 稀疏（MoE Top-k） | 稀疏（潜在MoE Top-k） → 输出： $\mathbf{h} \in \mathbb{R}^{B \times T \times D}$ → 骨干头（重复`L`次）

### 3.2 生成范式
图2 [https://arxiv.org/html/2609.17535#S3.F2] 展示了每种范式头如何安装在共享骨干输出 $\mathbf{h} \in \mathbb{R}^{B \times T \times D}$ 上，以及训练和推理在三种模式下的差异。

**自回归（AR）。** 在AR范式下，骨干采用标准因果注意力掩码 [Vaswani et al., 2017 (https://arxiv.org/html/2609.17535#bib.bib9)]，并通过最小化下一个token交叉熵进行优化。在推理阶段，通过KV缓存 [Shazeer, 2019 (https://arxiv.org/html/2609.17535#bib.bib45); Pope et al., 2023 (https://arxiv.org/html/2609.17535#bib.bib46)] 加速生成，以避免上下文重新计算。生成器支持不同的解码策略，包括贪心搜索、温度缩放、top-k [Fan et al., 2018 (https://arxiv.org/html/2609.17535#bib.bib43)] 和top-p采样 [Holtzman et al., 2019 (https://arxiv.org/html/2609.17535#bib.bib44)]。

**离散扩散。** 对于离散扩散，**DantinoX** 实现了LLaDA公式 [Nie et al., 2026 (https://arxiv.org/html/2609.17535#bib.bib25)]，建立在掩码扩散语言模型的基础之上 [Austin et al., 2021 (https://arxiv.org/html/2609.17535#bib.bib3); Sahoo et al., 2024 (https://arxiv.org/html/2609.17535#bib.bib56)]。在此模式下，骨干使用**双向**注意力。训练时，token根据可配置的时间步`t`调度（线性、余弦或平方根）[Austin et al., 2021 (https://arxiv.org/html/2609.17535#bib.bib3)] 被独立地替换为`[MASK]`token，并通过掩码位置上的$(1/t)$加权交叉熵损失进行优化。推理通过`S`个反向扩散步骤 [Ho et al., 2020 (https://arxiv.org/html/2609.17535#bib.bib81)] 执行，并内置支持多种解码策略（采样、贪心、置信度和因子）。该框架还提供了一个**块生成**模式 [Nie et al., 2026 (https://arxiv.org/html/2609.17535#bib.bib25); Wu et al., 2025 (https://arxiv.org/html/2609.17535#bib.bib55)]，该模式顺序地对局部窗口进行去噪；为了优化吞吐量，还提供了**双缓存** [Wu et al., 2025 (https://arxiv.org/html/2609.17535#bib.bib55)] 以在精度损失较低的情况下实现更大的加速。

**连续流匹配。** 对于连续流匹配，**DantinoX** 采用了最新的ELF公式 [Hu et al., 2026 (https://arxiv.org/html/2609.17535#bib.bib5)]。这种方法在冻结编码器（例如，T5 [Raffel et al., 2020 (https://arxiv.org/html/2609.17535#bib.bib52)]）的连续嵌入空间中操作，仅在最后去噪步骤通过共享权重网络进行离散化解码，从而避免了逐步的token级监督。训练时，骨干学习对应于前向过程 $\mathbf{z}_t = t \mathbf{x} + (1-t) \boldsymbol{\varepsilon}$（其中 $\mathbf{x}$ 是归一化嵌入，$\boldsymbol{\varepsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$，且 $t \sim U[0,1]$）的时间依赖向量场 [Lipman et al., 2022 (https://arxiv.org/html/2609.17535#bib.bib76)]。

**图2：生成范式及相应的推理策略。**  
骨干输出 $\mathbf{h} \in \mathbb{R}^{B \times T \times D}$ → 骨干头：  
- **离散（双向）**：掩码交叉熵（训练）→ 反向扩散 `S` 步（采样 / 置信度 / 因子）| 块（推理）  
- **自回归（因果）**：下一个token交叉熵（训练）→ KV缓存 / 无缓存 解码（贪心 / top-k / top-p）（推理）  
- **连续（双向）**：在 $\mathbb{R}^E$ 空间中的MSE + CE（训练）→ ODE / SDE `N` 步（CFG 尺度 `w`, 噪声 $\gamma$）（推理）