arXiv:2604.03537cs.CLcs.LG2026-04

用树结构降低语言模型预测复杂度,省内存还更高效

Rethinking Token Prediction: Tree-Structured Diffusion Language Model

  • 用预构建词表树分解预测任务,降低分类维度
  • 相同参数下显存减半,困惑度媲美顶尖模型
  • 适合资源受限场景,尤其利于小规模模型训练

离散扩散语言模型已成为自回归模型的有力替代,但在参数和内存受限条件下高效训练仍具挑战。现有架构普遍采用全词汇表令牌预测层,占模型参数比重较大(如小规模DiT设计中超过20%),且常主导峰值GPU内存使用,导致资源利用效率低下。为此,本文重新审视显式全词汇表预测的必要性,转而利用词元间的内在结构,构建树结构扩散语言模型。具体而言,将扩散过程中的中间隐状态与预构建词表树中词元的祖先节点对应,通过树结构因子分解实现分类维度的指数级降低,使预测头可忽略不计,并将节省的参数重新分配至加深注意力模块。实证表明,在相同参数预算下,该方法将峰值GPU内存使用量减少一半,同时保持与当前最先进离散扩散语言模型相当的困惑度表现。

原文摘要 · Abstract (English)

Discrete diffusion language models have emerged as a competitive alternative to auto-regressive language models, but training them efficiently under limited parameter and memory budgets remains challenging. Modern architectures are predominantly based on a full-vocabulary token prediction layer, which accounts for a substantial fraction of model parameters (e.g., more than 20% in small scale DiT-style designs) and often dominates peak GPU memory usage. This leads to inefficient use of both parameters and memory under constrained training resources. To address this issue, we revisit the necessity of explicit full-vocabulary prediction, and instead exploit the inherent structure among tokens to build a tree-structured diffusion language model. Specifically, we model the diffusion process with intermediate latent states corresponding to a token's ancestor nodes in a pre-constructed vocabulary tree. This tree-structured factorization exponentially reduces the classification dimensionality, makes the prediction head negligible in size, and enables reallocation of parameters to deepen the attention blocks. Empirically, under the same parameter budget, our method reduces peak GPU memory usage by half while matching the perplexity performance of state-of-the-art discrete diffusion language models.

扩散模型语言模型内存优化树结构

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。