lavanyanigam/soft-masking-checkpoints

Lost in Interpolation: Spherical Soft-Masked MDLM Checkpoints (5,000 Steps)

This model card contains checkpoints and information for the paper "Lost in Interpolation: Why Predictive Feedback Fails in Diffusion Language Models" (COLM 2026).

The checkpoints evaluate the effect of Spherical Soft-Masking (S-SM) using Spherical Linear Interpolation (SLERP) and Karcher flow against standard Euclidean linear interpolation (LERP) on a 169M-parameter Diffusion Transformer (DiT) backbone. All configurations were fine-tuned for 5,000 steps on the OpenWebText dataset.

Model Description

• Paper: Lost in Interpolation: Why Predictive Feedback Fails in Diffusion Language Models (COLM 2026 Under Review)
• Backbone Architecture: 169M-parameter Diffusion Transformer (DiT) pretrained for 1M steps (Sahoo et al., 2024)
• Tokenizer: GPT-2 BPE (vocab size: 50,258)
• Context Length: 1,024
• Objective: Investigate norm-collapse in soft-masking feedback. Rather than interpolating linearly (LERP) through Euclidean space $\mathbb{R}^D$ which leads to off-manifold token embeddings with systematically reduced norms, this approach maps feedback using Spherical Interpolation (SLERP) over the Fréchet mean of top-$k$ predictions to preserve the natural hyperspherical scale of token embeddings.

---

Evaluation & Results

Both interpolation schemes (LERP vs. SLERP) were evaluated under two regimes:

1. Fixed $\lambda = 0.3$: Prevents the optimizer from suppressing the feedback signal to isolate geometric differences.
2. Learned $\lambda$: Employs a learnable transparency head initialized near $\lambda \approx 0.13 - 0.14$.

Validation Metrics (OpenWebText validation split)

Feedback Configuration	val/BPD $\downarrow$	val/NLL $\downarrow$	val/PPL $\downarrow$
Vanilla MDLM (Sahoo et al., 2024)	4.5194	3.1326	22.9336
Fixed $\lambda = 0.3$ Regimes
├─ LERP SM	4.9719	3.4462	31.3821
└─ Our SLERP SM	4.7050	3.2614	26.0860
Learned $\lambda$ Regimes
├─ LERP SM	4.5300	3.1399	23.1030
└─ Our SLERP SM	4.5370	3.1450	23.2270

Unconstrained Generation Quality (5,000 samples)

Evaluated against OpenWebText reference text.

Feedback Configuration	Gen PPL $\downarrow$	Entropy $\uparrow$	MAUVE $\uparrow$
Vanilla MDLM	60.17	5.51	0.0201
Fixed $\lambda = 0.3$ Regimes
├─ LERP SM	48.25	5.52	0.0551
└─ Our SLERP SM (Ours)	26.91	5.35	0.2588
Learned $\lambda$ Regimes
├─ LERP SM	118.91	5.67	0.0071
└─ Our SLERP SM	51.38	5.49	0.0293

---

Method Details: Spherical Soft-Masking (S-SM)

For a masked position $l$ with a target top-$k$ prediction distribution $\pi$, Spherical Soft-Masking replaces standard Euclidean convex combinations with two geometric phases:

Phase A: Weighted Fréchet Mean via Karcher Flow

Finds the intrinsic center $\mu^\star$ on the sphere $\mathbb{S}^{D-1}$ by minimizing the weighted squared geodesic distances: $${\mu }^{\star }=\arg \underset{\mu \in {\mathbb{S}}^{D-1}}{\min }\sum _{i\in \text{top-}k(p)}{\pi }_i{d}_g(\mu ,{\widehat{\mathbf{v}}}_i{)}^2\backslash mu^\backslash star\; =\; \backslash arg\backslash min\_\{\backslash mu\; \backslash in\; \backslash mathbb\{S\}^\{D-1\}\}\; \backslash sum\_\{i\; \backslash in\; \backslash text\{top-\}k(p)\}\; \backslash pi\_i\; \backslash ,\; d\_g(\backslash mu,\; \backslash hat\{\backslash mathbf\{v\}\}\_i)^2$$ Since there is no closed-form solution, it is solved using Karcher fixed-point iteration (typically $N_{\text{iter}} = 3$ is sufficient): $$\mu \leftarrow {\text{exp}}_{\mu }(\sum _{i\in \text{top-}k}{\pi }_i{\log }_{\mu }({\widehat{\mathbf{v}}}_i))\backslash mu\; \backslash leftarrow\; \backslash text\{exp\}\_\{\backslash mu\}\; \backslash left(\; \backslash sum\_\{i\; \backslash in\; \backslash text\{top-\}k\}\; \backslash pi\_i\; \backslash log\_\backslash mu(\backslash hat\{\backslash mathbf\{v\}\}\_i)\; \backslash right)$$

Phase B: Geodesic SLERP Blend

Interpolates along the great-circle arc between the normalized mask embedding $\hat{\mathbf{m}}$ and the target mean $\mu^\star$: $$\mathbf{s}=\frac{\sin ((1-\lambda )\mathrm{\Omega })}{\sin \mathrm{\Omega }}\widehat{\mathbf{m}}+\frac{\sin (\lambda \mathrm{\Omega })}{\sin \mathrm{\Omega }}{\mu }^{\star },\text{where }\mathrm{\Omega }=\arccos ⟨\widehat{\mathbf{m}},{\mu }^{\star }⟩\backslash mathbf\{s\}\; =\; \backslash frac\{\backslash sin((1\; -\; \backslash lambda)\backslash Omega)\}\{\backslash sin\backslash Omega\}\backslash hat\{\backslash mathbf\{m\}\}\; +\; \backslash frac\{\backslash sin(\backslash lambda\backslash Omega)\}\{\backslash sin\backslash Omega\}\backslash mu^\backslash star,\; \backslash quad\; \backslash text\{where\; \}\; \backslash Omega\; =\; \backslash arccos\backslash langle\backslash hat\{\backslash mathbf\{m\}\},\; \backslash mu^\backslash star\backslash rangle$$

Finally, the unit direction $\mathbf{s}$ is rescaled to the native token norm to preserve backbone compatibility: $${\mathbf{x}}^{\text{sph}}=\mathrm{\parallel }\mathbf{m}\mathrm{\parallel }\mathbf{s}\backslash mathbf\{x\}^\{\backslash text\{sph\}\}\; =\; \backslash |\backslash mathbf\{m\}\backslash |\backslash ,\backslash mathbf\{s\}$$

---

Training Hyperparameters

• Backbone learning rate ($\eta_{bb}$): $3 \times 10^{-4}$
• SM parameter learning rate ($\eta_{sm}$): $10^{-2}$
• Top-$k$ width ($k$): $3$
• Karcher iterations ($N_{iter}$): $3$
• SM activation probability ($p_{sm}$): $0.8$
• Timestep band ($b_l / b_h$): $0.2 ,/, 0.8$
• Numerical tolerance ($\varepsilon$): $10^{-6}$

---