Lost in Interpolation: Spherical Soft-Masking (SLERP-SM) for Diffusion Language Models

Official PyTorch implementation of Spherical Soft-Masking (SLERP-SM) for masked diffusion language models (MDLMs), from the paper:

Lost in Interpolation: Why Predictive Feedback Fails in Diffusion Language Models (COLM 2026)

Soft-masking provides predictive feedback by continuously blending the mask token embedding with a superposition of top-k predictions. However, standard linear interpolation (LERP-SM) suffers from a severe norm-collapse failure:

• Hyperspherical Geometry: Language model token embeddings concentrate on a spherical shell in high dimensions. Straight-line interpolation (LERP) cuts through the hypersphere, systematically shrinking embedding norms by 20–40% and feeding out-of-distribution inputs to the backbone.
• Spherical Soft-Masking (SLERP-SM):
  1. Fréchet (Karcher) Mean: Aggregates top-$k$ normalized predictions on the unit sphere $\mathcal{S}^{d-1}$.
  2. Geodesic Interpolation (SLERP): Interpolates along the great-circle arc between the mask embedding and the Fréchet mean.
  3. Norm Preservation: Rescales the interpolated direction back to the original mask embedding norm, preventing norm collapse and gradient explosion.

This repository contains implementations for both Language Modeling and Code Generation:

Built on MDLM / Duo:

cd language
conda create -n sm-env python=3.12 -y && conda activate sm-env
pip install -r requirements.txt
pip install flash_attn==2.7.4.post1 --no-build-isolation

# Finetune with SLERP-SM (multi-GPU)
NUM_GPUS=2 BASE_CKPT=/path/to/mdlm.ckpt bash scripts/run_slerp_multigpu.sh

See language/README.md for full pretraining, evaluation, and Modal / Slurm scripts.

Built on Dream-7B with PEFT / DoRA:

cd coding/train
python train.py "./configs/base_configs/config_1.json" "./configs/datasets_softmasking/config_coding.json"

See coding/README.md for fine-tuning and HumanEval / EvalPlus evaluation.

Pretrained checkpoints for Language Modeling (OpenWebText) are available on Hugging Face: 👉 HuggingFace Checkpoints (lavanyanigam/soft-masking-checkpoints)

This codebase builds upon Duo, ReMDM, Dream-7B, and Soft-Masked DLMs.