Lost in Interpolation: Why Predictive Feedback Fails in Diffusion Language Models
Aryan Sood ⋅ Ishaan Bansal ⋅ Lavanya Nigam ⋅ Vidit Aggarwal
Abstract
Soft masking has emerged as an effective technique to accelerate the convergence of Masked Diffusion Language Models (MDLMs) during fine-tuning and continued pre-training. However, existing formulations implicitly assume that the embedding space is Euclidean, incorporating linear feedback from the model's top-$k$ token predictions. We empirically challenge this assumption: a geodesic angle diagnostic reveals that the mask token and predicted token embeddings subtend a mean angle of $\sim$73$^\circ$ throughout training and the norm of the embeddings remains near constant across all frequency ranks. These jointly indicate a hyperspherical geometry for which linear interpolation is inappropriate. We close this gap with a soft-masking framework governed by Spherical Interpolation. We instantiate this feedback mechanism via the Fr\'echet mean and Karcher flow, with an explicit norm-restoration step to maintain backbone compatibility. Our experiments on OpenWebText demonstrate that this geometrically-aware approach eliminates the training degradation suffered by standard (LERP-based) soft-masked MDLMs when trained with a fixed mixing weight ($\lambda$), converging on par with vanilla MDLM training rather than collapsing below it. Furthermore, analysis of the learned mixing weight($\lambda$) reveals that the optimizer actively suppresses linear feedback over time, halving its magnitude compared to our spherical formulation. Under a fixed mixing weight, our SLERP instantiation also yields markedly higher-quality generations, achieving a MAUVE score of 0.258 versus 0.055 for LERP.
Chat is not available.
Successful Page Load