PreDiff-LM: Pretrained Discrete Masked Diffusion Language Modeling with Hybrid Attention

Natural language generation has evolved from simple n-gram models to complex neural architectures like GPT-2 and BERT. Autoregressive models excel in left-to-right text prediction but struggle with infilling and bidirectional tasks. Discrete diffusion models, inspired by image generation techniques, introduce a corruption process that adds noise to sequences and learn to denoise iteratively. Early diffusion models like Austin et al. (2021) and Lou et al. (2023) demonstrated bidirectional capabilities but faced issues with training efficiency and quality. Recent efforts, such as DiffuGPT (Gong et al., 2024), have attempted to adapt pretrained autoregressive models, but attention pattern mismatches hinder optimal transfer. The challenge remains to reconcile causal pretraining with bidirectional denoising effectively.

The core problem addressed is how to adapt pretrained causal transformers for discrete diffusion tasks without losing their inherent strengths. The attention pattern mismatch—causal over prompts versus bidirectional over targets—limits the effectiveness of transfer. Fully replacing causal masks with bidirectional ones degrades prompt representations, while maintaining causality prevents full bidirectional denoising. This tension results in suboptimal generation quality and training inefficiency. Finding a mechanism that preserves the prompt’s causal structure while enabling full bidirectional context for targets is crucial for leveraging pretrained models in diffusion frameworks.

The main innovation is the asymmetric hybrid attention mask that selectively combines causal and bidirectional attention. This mask keeps prompt tokens under causal attention, ensuring the preservation of pretraining structures, while allowing target tokens to attend bidirectionally, facilitating effective denoising. Additional contributions include noise level embeddings, self-conditioning to improve denoising consistency, and confidence-aware unmasking to control decoding errors. These components work synergistically to enable efficient transfer from autoregressive to diffusion models, significantly reducing training steps and improving generation metrics. The approach represents a fundamental shift in attention design for diffusion-based language modeling.

• �� Develop an asymmetric attention mask partitioning input into prompt and target regions, with prompt tokens maintaining causal attention and target tokens attending bidirectionally. • Encode noise levels using sinusoidal functions plus an MLP, injecting this information into token representations to guide denoising. • Implement self-conditioning by feeding previous soft predictions back into the model, reducing inconsistencies among masked tokens. • During training, optimize a cross-entropy loss over masked positions with a noise schedule that emphasizes moderate noise levels. • In inference, utilize confidence scores to iteratively unmask tokens, starting from uncertain positions and gradually revealing high-confidence tokens. • Combine these mechanisms with a carefully designed noise schedule and hybrid attention to balance training efficiency and generation quality. • Validate through experiments on WikiText-103, comparing against baselines like pure causal, pure bidirectional, and uniform attention models.

The experimental setup involves training models of various sizes (from 133M to 1.56B parameters) on WikiText-103, with 90K steps and a batch size of 256. Baselines include original GPT-2, fine-tuned autoregressive controls, and multiple diffusion models such as MDLM, DiffuGPT, and SEDD. Metrics encompass perplexity, MAUVE, repetition rate, diversity, and zero-shot task performance. Ablation studies test the importance of attention masks, noise scheduling, and self-conditioning. Results show that the hybrid attention mechanism reduces perplexity significantly and accelerates training convergence, with models reaching perplexity below 50 in just 8K steps, compared to 350K for previous methods.

The hybrid attention model achieves a perplexity of 28.7 on WikiText-103, outperforming the uniform attention baseline at 34.1. MAUVE scores improve from 0.71 to 0.78, indicating better distributional quality. Pretraining reduces training steps to reach perplexity 50 from 350K to 8K, demonstrating high efficiency. Combining with DiffuGPT objectives further lowers perplexity to 26.9. The model also exhibits reduced repetition, improved diversity, and better human preference ratings, confirming the effectiveness of the attention adaptation. Performance scales favorably with model size, narrowing the gap with autoregressive controls.

This approach benefits tasks requiring flexible text generation, such as content creation, dialogue systems, and infilling applications. Its ability to generate high-quality, diverse outputs with fewer steps makes it suitable for deployment in resource-constrained environments. The method can be extended to multi-modal tasks, including image captioning and video synthesis, by integrating visual features. Its efficiency and adaptability also support real-time applications like chatbots and AI assistants, where balancing quality and speed is critical.

Despite promising results, the model's inference speed at higher denoising steps remains slower than optimized autoregressive decoders, limiting real-time use. The approach relies heavily on high-quality pretrained weights, and its effectiveness may diminish with less suitable architectures or domain shifts. Additionally, the experiments are primarily on WikiText-103, and broader validation across diverse datasets and tasks is needed. Future work should address robustness in high-noise regimes, improve decoding speed, and explore multi-modal extensions to enhance practical deployment.