MDLMPE:面向掩码扩散语言模型的分布感知位置编码

MDLMPE: Distribution Aware Positional Encoding for Masked Diffusion Language Models

精选理由

掩码扩散模型的位置编码一直照搬自回归那套,MDLMPE 第一个把遮盖状态算进去,在 LLaDA 和 DREAM 上效果更好。

AI 摘要

MDLMPE 将掩码扩散模型中 token 的可见/掩码状态编码为二进制序列,结合高斯权重与余弦基投影生成位置特征,再经轻量 MLP 映射为角度偏移来调制 RoPE 相位。在 LLaDA 和 DREAM 上的实验覆盖监督微调、预训练、零样本评测和块扩散设置,整体优于传统位置编码。消融实验表明,可用性状态、高斯局部性、频谱基与嵌入注入的组合效果最强。

原文 · arXiv cs.AI

MDLMPE: Distribution Aware Positional Encoding for Masked Diffusion Language Models

Masked diffusion language models (MDLMs) enable parallel generation and bidirectional context modeling, but their positional context differs fundamentally from that of autoregressive (AR) models. Whereas AR decoding exposes a contiguous prefix, MDLM denoising produces dynamic, non-contiguous configurations of revealed and masked tokens. Conventional positional encodings such as RoPE capture sequence order and pairwise displacement but remain insensitive to this evolving token-availability structure. To address this limitation, we propose MDLMPE, a positional encoding designed specifically for masked diffusion. To the best of our knowledge, MDLMPE is the first method to make positional representations explicitly aware of the changing revealed/masked configuration. It represents token availability as a binary sequence, applies distance-aware Gaussian weighting, and projects the resulting pattern through a cosine basis to obtain distribution-aware positional features. These features are added to token embeddings and mapped by a lightweight MLP to angular offsets that modulate the standard RoPE phases. Extensive experiments on LLaDA and DREAM demonstrate that MDLMPE generally outperforms conventional positional encoding methods across supervised fine-tuning, pretraining, zero-shot evaluation, and block-diffusion settings. Further ablations show that the complete combination of availability state, Gaussian locality, spectral basis, and embedding injection yields the strongest result. These results establish the evolving token-availability distribution as a useful positional signal for masked diffusion language models.