让生成的错词能被重新掩码修正,提升文本生成质量。
Don't Settle Too Early: Self-Reflective Remasking for Diffusion Language Models
- 引入重掩码机制,动态识别并重掩低质量词
- 在多个数据集上达到开源扩散语言模型最佳效果
- 适合需要高质量文本生成的研究与应用
基于掩码的扩散语言模型(DLMs)难以修正已生成的错误词汇:一旦生成,通常无法更改。核心挑战在于识别输入中的潜在错误。本文提出一种名为RemeDi的掩码式扩散语言模型,引入「重掩码」作为基础机制,实现更灵活的文本精炼。RemeDi在每一步同时预测词分布与每个词的置信度分数,根据置信度决定哪些词应被解掩码以重新生成,从而识别低质量词并重新采样。为此,我们设计了包含监督微调和强化学习的训练流程:前者教会模型检测并重掩错误词,后者优化完整生成轨迹以获得更高奖励。实验表明,RemeDi在多个数据集上超越现有开源DLMs,达到当前最优水平。
原文摘要 · Abstract (English)
Mask-based Diffusion Language Models (DLMs) struggle to revise incorrect tokens: once a token is generated, it typically remains fixed. The key challenge is to identify potential errors in the inputs. In this paper, we propose \emph{\underline{Rem}asking-\underline{e}nabled \underline{Di}ffusion Language Model (RemeDi}, a mask-based DLM that introduces \emph{remasking} as another fundamental mechanism, enabling more flexible text refinement in diffusion-based text generation. To achieve this, RemeDi jointly predicts token distributions and per-token confidence scores at each step. The confidence scores determine which tokens to be unmasked after the current step, allowing the model to identify tokens with low quality and remask them. These remasked tokens can be resampled with richer context in subsequent steps. We design a remask-aware pipeline to train this ability, including supervised fine-tuning which teaches the model to detect and remask incorrect tokens in addition to predict mask tokens, and reinforcement learning which optimizes full generation trajectories toward higher rewards. Experiments show that RemeDi achieves the state-of-the-art results among open-source DLMs on multiple datasets.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。