---
title: Masked Prediction Distribution (MPD)
url: https://www.emergentmind.com/topics/masked-prediction-distribution-mpd
type: topic
---

# Masked Prediction Distribution (MPD)

Searching arXiv for the cited papers to ground the article in current sources.
Masked Prediction Distribution (MPD) denotes the conditional law used to predict masked or skipped content from partially observed context, but its precise formalization depends on the modeling regime. In causal audio generation with continuous-valued tokens, MPD is the diffusion-based conditional distribution $p_\theta(x_f \mid x_{< i}, w, v, p_t)$ for a future token selected by a masking pattern [2507.09834]. In masked diffusion language models, MPD can be defined as the per-position categorical distribution $p_\theta(\cdot \mid X_{\text{in}}(t))^i$ induced by the current continuous decoding state [2606.29066]. In the identifiability analysis of masked prediction tasks, MPD is the true conditional law $P(x_t \mid x_{-t}; \theta)$ under the generative model, and the optimal masked predictor should realize it [2202.09305]. Taken together, these formulations place MPD at the intersection of conditional modeling, self-supervised prediction, and decoding dynamics.

## 1. Definitions across modeling paradigms

The cited literature does not present MPD as a single universal object; instead, it specifies closely related conditional distributions for different architectures and data types. A concise comparison is useful.

| Setting | MPD form | Role |
|---|---|---|
| Causal audio LM | $p_\theta(x_f \mid x_{< i}, w, v, p_t)$ | Diffusion-based next-token distribution |
| MDLM continuous decoding | $p_\theta(\cdot \mid X_{\text{in}}(t))^i$ | Per-step belief over tokens |
| Masked-prediction theory | $P(x_t \mid x_{-t}; \theta)$ | True conditional law under the model |

In "Generative Audio Language Modeling with Continuous-valued Tokens and Masked Next-Token Prediction" [2507.09834], the relevant task is masked next-token prediction with continuous-valued tokens. An audio clip is mapped by a VAE encoder into a sequence of continuous tokens $x = \{x_1,\dots,x_n\}$ with $x_t \in \mathbb{R}^h$, conditioned on a text prompt $w$. A mask variable $v = \{v_1,\dots,v_n\}$ drops tokens to form a shorter sequence $x^v$, and for a visible position $i$ the prediction target is $x_{\mathrm{succ}_v(i)}$, a future token that may be multiple steps ahead in original time. The MPD is then the conditional next-token distribution for that target under causal constraints [2507.09834].

In "Masked Diffusion Decoding as $x$-Prediction Flow" [2606.29066], MPD is introduced as an inferred term for the per-position predictive categorical law in masked diffusion language models. The vocabulary is augmented with a special mask token $[M]$, and the model produces logits $z_{\text{pred}}^i = f_\theta(\tilde{x})^i \in \mathbb{R}^{|\mathcal{V}|}$ and probabilities $p_\theta(\cdot \mid \tilde{x})^i = \operatorname{softmax}(z_{\text{pred}}^i)$. Under continuous decoding, the same object is evaluated on a continuous input state $X_{\text{in}}(t)$ rather than only on a discrete masked string, yielding a continuously evolving MPD [2606.29066].

In "Masked prediction tasks: a parameter identifiability view" [2202.09305], MPD is the true conditional law of a masked token given the observed context under a parametric latent-variable model. For a masked position $t$, the quantity of interest is $P(x_t \mid x_{-t}; \theta)$. Under proper loss functions, the optimal masked predictor recovers this true conditional; under squared loss it becomes the conditional expectation, and for discrete one-hot targets that conditional expectation equals the categorical MPD vector [2202.09305].

## 2. MPD in causal audio language modeling with continuous-valued tokens

The audio formulation is explicitly probabilistic and diffusion-based. For any visible position $i$ and target $f = \mathrm{succ}_v(i)$, the masked prediction distribution is
$$
p_\theta(x_f \mid x_{< i}, w, v, p_t).
$$
The Transformer decoder is causal: it sees the visible past $x_{< i}$, the prompt $w$, a BOS token, and a target positional embedding $p_t$ indicating which future index is to be predicted. The diffusion head then defines the MPD through a token-wise denoising diffusion process [2507.09834].

The forward process uses the standard variance-preserving DDPM parameterization with a cosine schedule. With $\{\beta_k\}_{k=1}^K$, $\alpha_k = 1-\beta_k$, and $\bar{\alpha}_k = \prod_{s=1}^k \alpha_s$, the noised target token satisfies
$$
q(x_f^{(k)} \mid x_f) = \mathcal{N}\!\left(\sqrt{\bar{\alpha}_k}\, x_f,\; (1-\bar{\alpha}_k) I\right),
$$
equivalently,
$$
x_f^{(k)} = \sqrt{\bar{\alpha}_k}\, x_f + \sqrt{1-\bar{\alpha}_k}\,\epsilon,\qquad \epsilon \sim \mathcal{N}(0,I).
$$
The reverse process conditions on the causal context vector $z_i$ produced by the Transformer decoder and the diffusion step $k$. With an $\epsilon$-prediction parameterization, the diffusion head predicts
$$
\hat{\epsilon}_\theta = M_\phi(x_f^{(k)}, z_i, k, p_t),
$$
and the reverse transition is Gaussian,
$$
p_\theta(x_f^{(k-1)} \mid x_f^{(k)}, z_i) = \mathcal{N}\!\left(\mu_\theta(x_f^{(k)}, z_i, k),\; \sigma_k^2 I\right),
$$
with mean
$$
\mu_\theta(x_f^{(k)}, z_i, k)
=
\frac{1}{\sqrt{\alpha_k}}
\left[
x_f^{(k)} - \frac{\beta_k}{\sqrt{1-\bar{\alpha}_k}}\, \hat{\epsilon}_\theta
\right].
$$
This yields the implicit diffusion-chain definition
$$
p_\theta(x_f \mid x_{< i}, w, v, p_t)
=
\int p_\theta(x_f^{(0:K)} \mid x_{< i}, w, v, p_t)\, dx_f^{(1:K)},
$$
with prior $p(x_f^{(K)}) = \mathcal{N}(0,I)$ [2507.09834].

Training minimizes the standard denoising diffusion MSE loss token-wise, averaged over dataset samples $(w,x)$, masking patterns $v \sim P_{\text{mask}}$, visible positions $i \in V \setminus \{\max(V)\}$, diffusion steps $k \sim U\{1,\dots,K\}$, and Gaussian noise $\epsilon \sim \mathcal{N}(0,I)$:
$$
L(\theta,\phi)
=
\mathbb{E}
\left[
\|\epsilon - M_\phi(x_f^{(k)}, z_i, k, p_t)\|_2
\right].
$$
The paper emphasizes that masking is implemented by dropping rather than inserting special mask tokens, reducing sequence length and compute when the masking ratio is high. The masking-ratio distribution is a mixture schedule,
$$
P_{\text{mask}}(r) = 0.5\, \text{Normal}(\mu,\sigma^2)\text{ on }[0,1] + 0.5\, \text{TruncNormal}(a,b)\text{ on }[0,1],
$$
where the Normal component emphasizes high masking ratios and the truncated-normal component contributes a long tail of low ratios to reduce train-test mismatch for standard next-token decoding [2507.09834].

This formulation differs from both standard autoregressive next-token prediction on discrete tokens and masked language modeling. In the discrete autoregressive setting, $p(x_t \mid x_{< t})$ is modeled by a softmax over a fixed vocabulary and trained with cross-entropy. Here the next-token law is continuous and modeled via token-wise diffusion, with denoising MSE on $\epsilon$ rather than cross-entropy. In contrast to bidirectional MLM, the masked next-token task is strictly causal at train and test time: the model predicts a future token using only the visible past, and the target positional embedding disambiguates which future index is being predicted. The paper states that removing this target positional embedding severely hurts performance because different future targets otherwise conflict during training [2507.09834].

## 3. MPD in masked diffusion language models and $x$-prediction flow

In masked diffusion language models, MPD is the token distribution the decoder predicts at each position and each diffusion-progress value. The basic MDLM training objective is masked cross-entropy over masked positions,
$$
\mathcal{L}_{\text{MDLM}}(\theta)
=
-
\mathbb{E}_{\sigma, x_0, \tilde{x}}
\left[
\frac{1}{\sigma}
\sum_{i=1}^{N}
\mathbf{1}[i \in \mathcal{M}]
\log p_\theta(x_0^i \mid \tilde{x})
\right],
$$
and for an input $\tilde{x}$ the mask predictor produces logits and per-position probabilities
$$
z_{\text{pred}}^i = f_\theta(\tilde{x})^i \in \mathbb{R}^{|\mathcal{V}|},
\qquad
p_\theta(\cdot \mid \tilde{x})^i = \operatorname{softmax}(z_{\text{pred}}^i).
$$
The MPD at position $i$ and progress $t$ is then defined as
$$
\mathrm{MPD}_i(t) \coloneqq p_\theta(\cdot \mid X_{\text{in}}(t))^i = \operatorname{softmax}(z_{\text{pred}}^i(X_{\text{in}}(t))).
$$
This is the central object preserved across continuous decoding steps [2606.29066].

The continuous state is anchored at the mask embedding $m$ and interpolates toward the clean embedding $X_0$ by a per-token progress vector $t \in [0,1]^N$:
$$
X_{\text{in}} = t \cdot X_0 + (1-t)\cdot m.
$$
At $t=0$, the state is the mask embedding; at $t=1$, it is the clean embedding. The model forms a clean-state prediction either by an argmax-embedding readout,
$$
X_{\text{pred}} = \operatorname{Embed}(\operatorname{argmax}(f_\theta(X_{\text{in}}))),
$$
or by a soft embedding readout,
$$
X_{\text{pred}} = \operatorname{softmax}(Z_{\text{pred}})\cdot W_{\text{in}},
\qquad
Z_{\text{pred}} = f_\theta(X_{\text{in}}).
$$
A velocity field then moves the current state toward the predicted clean state,
$$
V(X_{\text{in}}, t) = \frac{X_{\text{pred}} - X_{\text{in}}}{1-t},
$$
with update
$$
X_{\text{in}} = X_{\text{in}} + V \cdot \Delta t.
$$
Because $\mathrm{MPD}_i(t)$ is evaluated from $X_{\text{in}}(t)$ at every step, its evolution is induced by this state update; the paper states that the distribution sharpens as $X_{\text{in}}(t)$ approaches the clean embedding [2606.29066].

A further alignment objective stabilizes clean-state prediction in embedding space:
$$
X_{\text{pred}} = \operatorname{softmax}(Z_{\text{pred}})\cdot W_{\text{in}},
\qquad
\mathcal{L}_{x\text{-pred}} = \|X_{\text{pred}} - X_0\|_2^2.
$$
This alignment is introduced because continuous decoding requires stability on clean inputs as well as masked ones [2606.29066].

The decoding schedule is asynchronous and confidence-based. Each token has its own progress $t_i$, updated by
$$
\Delta t = a \cdot (1-t),
\qquad
t = t + \Delta t,
$$
where $a \in [a_{\min}, a_{\max}]^N$ is selected by a learned policy. Confidence at position $i$ is
$$
c_i = \max(\operatorname{softmax}(z_i)).
$$
The paper also adds two discrete adjustments on top of the continuous flow. First, re-editing resets a token if its confidence falls substantially below its current commitment:
$$
x_i = c_i \cdot x_{\text{pred}} + (1-c_i)\cdot m,
\qquad
t_i = c_i,
$$
when $c_i - t_i < -0.1$. Second, hard commitment sets the token with the highest confidence among those with $t_i < 0.99$ to $t_i = 1$ at each step [2606.29066].

The paper’s interpretation is explicit: standard discrete masked decoding discards the predictive distribution between steps by reducing it to a committed token or a fully masked state, whereas continuous decoding preserves and exploits the full MPD. This enables revisable partial progress, token-wise asynchrony, and more efficient use of limited decoding budget [2606.29066].

## 4. MPD as the optimal conditional law in masked-prediction theory

The theoretical treatment in the identifiability paper starts from a generative model rather than a decoder. There, MPD is the true conditional law of a masked token given observed context:
$$
P(x_t \mid x_{-t}; \theta).
$$
The principal question is whether this object, or a family of such objects induced by a masked prediction task, identifies the parameters of the underlying latent-variable model. The paper studies Hidden Markov Models with both discrete and conditionally Gaussian observations and defines identifiability as injectivity of the mapping from the model parameters, modulo hidden-state permutations, to the optimal predictors [2202.09305].

In the discrete HMM case, with transition matrix $mTrans \in \mathbb{R}^{k \times k}$ and emission matrix $mEmiss \in \mathbb{R}^{d \times k}$, the optimal masked predictor for $x_2$ given $x_1=x$ is
$$
f^{2 \mid 1}(x) = \mathbb{E}[x_2 \mid x_1 = x] = mEmiss\, mTrans\, \phi(x),
$$
where $\phi(x)$ is the posterior over the hidden state at time $1$ given $x_1=x$. Since $x_2$ is a one-hot vector, this conditional expectation is exactly the categorical MPD vector $P(x_2 \mid x_1=x)$. The paper also gives the more explicit one-sided formula
$$
P(X_2 \mid X_1=x) = B A \phi(x)
$$
under its notation and assumptions [2202.09305].

For conditionally Gaussian HMMs with identity covariance and emission means $\mu_s \in \mathbb{R}^d$, the MPD for $x_2$ given $x_1=x$ is a mixture of Gaussians,
$$
P(x_2 \mid x_1=x)
=
\sum_{s=1}^k w_s(x)\, \mathcal{N}(\mu_s, I),
$$
where $w(x) = mTrans\, \phi(x)$. Its conditional mean, which is the optimal square-loss predictor, is
$$
\mathbb{E}[x_2 \mid x_1=x] = mMeans\, mTrans\, \phi(x),
$$
with $mMeans = [\mu_1,\dots,\mu_k]$ [2202.09305].

For general masked positions in an HMM with both-side context, the MPD is obtained by smoothing. Let $A = mTrans$ and $B = mEmiss$. Then the latent posterior at the masked time is
$$
P(S_t=s \mid X_{-t}) \propto \alpha_t(s)\beta_t(s),
$$
with standard forward-backward recursions, and the masked-token law is
$$
P(X_t=x \mid X_{-t}) = \sum_{s=1}^k P(S_t=s \mid X_{-t})\, B_{x,s}.
$$
The optimal predictor for one-hot $x_t$ equals this MPD vector [2202.09305].

The paper’s central conclusion is that the informativeness of MPD depends strongly on the masked task. For discrete HMMs, pairwise single-token predictions such as $x_2 \mid x_1$ are non-identifiable, even when multiple pairwise tasks are combined, because they only constrain non-unique matrix products. By contrast, multi-token masked prediction of adjacent tokens, such as $x_2 \otimes x_3 \mid x_1$, is identifiable via tensor decomposition under Kruskal-rank conditions. For Gaussian HMMs, even the pairwise task $x_2 \mid x_1$ is identifiable, due to the structure of the posterior $\phi(x)$ and the identifiability of Gaussian mixtures [2202.09305].

## 5. Architectures, objectives, and decoding mechanisms

Although the three treatments of MPD arise in different settings, each couples a conditional distribution with a specific optimization and inference mechanism.

In the audio model, MPD is realized by a decoder-only Transformer $C_\theta$ plus a diffusion head $M_\phi$. The backbone takes text prompt tokens from CLAP and FLAN-T5 embeddings, a BOS token, visible past audio tokens, content positional embeddings $p_c$, and a target positional embedding $p_t$. The diffusion head is a small MLP predicting $\hat{\epsilon}_\theta = M_\phi(x_f^{(k)}, z_i, k, p_t)$. The paper uses a cosine schedule, $K_{\text{train}} = 1000$ diffusion steps and $K_{\text{sample}} = 100$ steps by default. It states that the MLP diffusion head is the only change from a standard decoder LLM head, and that all other infrastructure, including KV-cache and streaming, remains applicable [2507.09834].

Inference in that system is standard left-to-right decoding with no masking and $p_t$ set to the next position, $p_t = p_c(i)+1$. Each next token is sampled by a diffusion denoising chain from Gaussian noise to the clean token. The paper also reports optional classifier-free guidance: during training, the prompt $w$ is replaced with a learned “fake/uncond” embedding $w_f$ with probability $10\%$, and at test time the conditional and unconditional predictions are blended with an annealed guidance scale. It states that $\omega_0 \approx 7$ is a good default and standard temperature $T=1$ is used [2507.09834].

In the masked diffusion decoding paper, the backbone MDLM is complemented by a lightweight policy network that chooses per-token step sizes from features extracted from the MPD and decoding state. The policy feature vector is
$$
s = [p_{(1)}, p_{(2)}, p_{(3)}, p_{(4)}, H_k, p_m, t, \rho],
$$
where $p_{(j)}$ are the top probabilities, $p_m = p_{(1)}-p_{(2)}$ is the confidence margin, $H_k$ is normalized entropy over the top-$k$ probabilities with $k=4$, $t$ is current token progress, and $\rho = \ell/(T-1)$ is normalized global step index. The policy is a two-layer MLP with SiLU activations and predicts Beta mean and concentration,
$$
\mu_\theta(s) = 0.05 + 0.9 \cdot \sigma(W_\mu h_\theta(s) + b_\mu),
$$
$$
\kappa_\theta(s) = \kappa_{\min} + (\kappa_{\max}-\kappa_{\min})\cdot \sigma(W_\kappa h_\theta(s) + b_\kappa),
$$
followed by
$$
y \sim \operatorname{Beta}(\alpha_\theta,\beta_\theta),
\quad
\alpha_\theta = \mu_\theta(s)\kappa_\theta(s),
\quad
\beta_\theta = (1-\mu_\theta(s))\kappa_\theta(s),
$$
and
$$
a = 2^u,
\qquad
u = \log_2 a_{\min} + y(\log_2 a_{\max}-\log_2 a_{\min}).
$$
The policy is trained with GRPO, and the reward is
$$
R_{b,g,i} = R^{\text{task}}_{b,g} + \lambda \cdot t^{\text{final}}_{b,g,i}.
$$
This makes MPD not only a prediction object but also a control signal for scheduling and commitment decisions [2606.29066].

In the identifiability framework, the “architecture” is deliberately abstracted away. The critical assumption is that the predictor class contains the true conditional law, so that under correct specification, infinite data, and proper loss, the minimizer equals the MPD. This suggests a methodological distinction: in the generative audio and MDLM papers, MPD is operationalized by specific neural parameterizations and sampling rules, whereas in the identifiability paper it is the target statistical object against which predictor classes are evaluated [2202.09305].

## 6. Empirical behavior, misconceptions, and open directions

The audio paper reports that continuous tokens plus diffusion next-token prediction already outperform a previous discrete autoregressive baseline. On AudioCaps, using only AudioCaps plus WavCaps for training, AudioNTP Base (193M) yields $20\%$ and $40\%$ relative improvements over AudioGen Base (285M) in FAD and KL respectively. Adding masked next-token prediction to obtain AudioMNTP Base yields a further $26\%$ FAD, $10\%$ KL, and $9\%$ CLAP relative gains over AudioNTP Base. AudioMNTP Large (462M) achieves the best FD and FAD among all compared systems and is near the leading diffusion models on KL, IS, and CLAP, while remaining causal and streamable. Relative FAD improvements versus AudioGen are $41\%$ against AudioGen Base (285M) and $33\%$ against AudioGen Large (1B). Human evaluation indicates that AudioMNTP approaches Tango 2 and significantly outperforms AudioGen in both text relevance and overall quality, with particular strength on speech OVL [2507.09834].

The masked diffusion decoding paper evaluates efficiency under limited decoding budget. On HumanEval with LLaDA-8B-Instruct, standard masked decoding at full budget $(T=L)$ obtains $46.34\%$, while $x$-prediction flow at $1/4$ budget obtains $45.12\%$, which the paper describes as $97\%$ of the full-budget baseline. On HumanEval with LLaDA2.0-mini, mask prediction at $1/4$ budget gives $32.32\%$ and $x$-prediction flow at the same budget gives $59.76\%$, a $+27.44$ absolute improvement. MBPP shows similar gains at $1/4$ budget: $+11.80$ for LLaDA-8B and $+12.40$ for LLaDA2.0-mini. Ablations at $1/4$ budget on HumanEval for LLaDA-8B report $45.12$ for full $x$-prediction flow, $41.46$ without hard commitment, $26.83$ without re-editing, and $17.07$ without $x$-prediction alignment [2606.29066].

Several common misconceptions are directly addressed by these results and formulations. First, masked prediction need not be bidirectional. The audio MNTP setup is strictly causal at train and test time, even though it incorporates masking by dropping tokens and predicting skip tokens from sparse past context [2507.09834]. Second, MPD is not restricted to discrete vocabularies. In the audio setting it is a continuous conditional distribution over latent vectors modeled by token-wise diffusion rather than by a softmax [2507.09834]. Third, preserving the full predictive distribution can matter independently of the underlying backbone. The MDLM paper argues that standard masked decoding discards runner-up probabilities and forces premature, irrevocable commitments, whereas continuous decoding retains the full MPD and exploits it for re-editing and asynchronous scheduling [2606.29066]. Fourth, stronger masked objectives do not automatically imply stronger parameter recovery. The identifiability results show that for discrete HMMs, even multiple pairwise MPDs can remain non-identifying, while adjacent multi-token prediction can become identifiable through tensor structure [2202.09305].

Open directions are stated most explicitly in the audio and MDLM work. The audio paper notes that a causal decoder remains less expressive than large, fully bidirectional latent diffusion systems, that the system currently depends on a VAE and a vocoder in the AudioLDM pipeline, and that parallel decoding via target positional manipulation is a promising research direction [2507.09834]. The MDLM paper identifies failure modes tied to the embedding-linearity assumption, premature or delayed commitments, oscillations caused by re-editing, and uneven contextual constraints under asynchronous schedules [2606.29066]. A plausible implication is that MPD is becoming not only a target of prediction but also an internal state variable for generation, control, and theoretical analysis.

Source: https://www.emergentmind.com/topics/masked-prediction-distribution-mpd