PC-Sampler: Position-Aware Confidence Calibration

• The paper introduces PC-Sampler, combining global trajectory planning with calibrated confidence to address biases in token selection for masked diffusion models.
• It utilizes a position-aware weight and a corpus-frequency penalty to suppress trivial tokens, yielding over 10% average improvement versus uncertainty-based samplers.
• Experimental evaluations across seven benchmarks demonstrate significant gains in reasoning and structural tasks compared to traditional MDM decoding methods.

Searching arXiv for the specified paper and closely related work to ground the article. Position-Aware Confidence-Calibrated Sampling (PC-Sampler) is an inference-time decoding strategy for masked diffusion LLMs that combines global trajectory planning with content-aware informativeness maximization (Huang et al., 18 Aug 2025). It was introduced to address two failure modes of widely adopted uncertainty-based samplers in masked diffusion models (MDMs): a lack of global trajectory control and a pronounced bias toward trivial tokens in the early stages of decoding (Huang et al., 18 Aug 2025). PC-Sampler incorporates a position-aware weighting mechanism to regulate the decoding path and a calibrated confidence score to suppress the premature selection of trivial tokens, and experiments on three advanced MDMs across seven benchmarks report that it outperforms existing MDM decoding strategies by more than 10% on average, significantly narrowing the performance gap with state-of-the-art autoregressive models (Huang et al., 18 Aug 2025).

1. Problem setting and motivation

Masked diffusion models are discrete sequence generators trained to recover original tokens from partially masked inputs. Given a clean sequence $x_0 = (x_0^1, \dots, x_0^L)$, training draws a time $t \sim \mathrm{Uniform}[0,1]$, independently masks each token with probability $t$ to obtain $x_t$, and minimizes

$$\mathcal{L}_{\mathrm{MDM}} = - \mathbb{E}_{t,x_0,x_t}\left[ \frac{1}{t}\sum_{i=1}^L \mathbf{1}[x_t^i = \mathbf{M}] \log p_\theta(x_0^i \mid x_t) \right].$$

At inference, decoding starts from a fully masked sequence and iteratively chooses one or more masked positions to unmask until all masks are resolved (Huang et al., 18 Aug 2025).

This freedom of unmasking order is both the advantage and the difficulty of MDMs. Unlike autoregressive models, MDMs can reveal tokens in any order, so generation quality is highly sensitive to decoding strategy. The paper reports that the same model with different samplers can show $>10\%$–$70\%$ swings in accuracy on reasoning benchmarks, which makes the sampler a first-order design choice rather than a minor implementation detail (Huang et al., 18 Aug 2025).

The baseline family examined by the paper consists of uncertainty-based samplers. At diffusion step $t$, for each masked position $i \in \mathcal{M}_t$, these methods compute a scalar score

$$s_t^i = \mathcal{F}\bigl(p_\theta(x^i \mid x_t)\bigr)$$

and greedily choose

$t \sim \mathrm{Uniform}[0,1]$0

The three common scoring functions discussed are confidence,

$t \sim \mathrm{Uniform}[0,1]$1

entropy-based certainty,

$t \sim \mathrm{Uniform}[0,1]$2

and margin,

$t \sim \mathrm{Uniform}[0,1]$3

where $t \sim \mathrm{Uniform}[0,1]$4 and $t \sim \mathrm{Uniform}[0,1]$5 are the top two tokens by probability (Huang et al., 18 Aug 2025).

The paper identifies two limitations of this family. First, uncertainty-based decoding produces a characteristic U-shaped trajectory: left and right boundary positions are decoded early, then the sampler fills inward. This occurs because boundary tokens such as start and end markers, newlines, and template fragments are highly predictable and therefore dominate local confidence scores (Huang et al., 18 Aug 2025). Second, these samplers exhibit a strong trivial token bias. In the first several diffusion steps, more than $t \sim \mathrm{Uniform}[0,1]$6 of selected tokens are reported to come from a hand-defined trivial set, including tokens such as <[EOS](https://www.emergentmind.com/topics/electro-optic-sampling-eos)>, newline, <SPACE>, <EOT>, punctuation, and filler phrases like “answer is” (Huang et al., 18 Aug 2025). The result is early commitment to structurally easy but semantically uninformative content.

2. Formal definition and decoding rule

PC-Sampler replaces purely local uncertainty scoring with a product of a positional prior and a calibrated confidence term: $t \sim \mathrm{Uniform}[0,1]$7 Here $t \sim \mathrm{Uniform}[0,1]$8 is a position-aware weight and $t \sim \mathrm{Uniform}[0,1]$9 is a calibrated confidence score that combines model confidence with a corpus-frequency penalty (Huang et al., 18 Aug 2025).

For a masked position $t$0, let the candidate token be the model’s top prediction

$t$1

and write

$t$2

The calibrated confidence is then defined as

$t$3

followed by clipping

$t$4

where $t$5 is an empirical token frequency estimated from a large background corpus $t$6, and $t$7 prevents extreme values for very rare tokens (Huang et al., 18 Aug 2025).

The decoding loop remains greedy at the position-selection level. At each step, the sampler computes $t$8 for every masked position, chooses

$t$9

samples

$x_t$0

updates the partially decoded sequence, and repeats until the mask set is empty (Huang et al., 18 Aug 2025). The method is therefore plug-and-play: it changes only the inference-time selection rule and requires no architectural modification or additional training.

The following decomposition summarizes the method’s two control variables.

Component	Definition	Role
Position-aware weight	$x_t$1	Global trajectory control
Calibrated confidence	$x_t$2, clipped by $x_t$3	Trivial-token suppression
Final score	$x_t$4	Position selection

This formulation is notable because the position prior is static while the calibrated confidence is dynamic. The global plan is encoded by $x_t$5, whereas the per-step adaptation comes from $x_t$6, which changes as the partial sequence acquires more context.

3. Position-aware trajectory control

The position-aware term is

$x_t$7

with position index $x_t$8 starting from $x_t$9 and hyperparameter $$\mathcal{L}_{\mathrm{MDM}} = - \mathbb{E}_{t,x_0,x_t}\left[ \frac{1}{t}\sum_{i=1}^L \mathbf{1}[x_t^i = \mathbf{M}] \log p_\theta(x_0^i \mid x_t) \right].$$0 (Huang et al., 18 Aug 2025). For $$\mathcal{L}_{\mathrm{MDM}} = - \mathbb{E}_{t,x_0,x_t}\left[ \frac{1}{t}\sum_{i=1}^L \mathbf{1}[x_t^i = \mathbf{M}] \log p_\theta(x_0^i \mid x_t) \right].$$1, earlier positions receive larger weights, inducing a soft left-to-right bias. For $$\mathcal{L}_{\mathrm{MDM}} = - \mathbb{E}_{t,x_0,x_t}\left[ \frac{1}{t}\sum_{i=1}^L \mathbf{1}[x_t^i = \mathbf{M}] \log p_\theta(x_0^i \mid x_t) \right].$$2, all positions are weighted equally, and the trajectory is driven entirely by calibrated confidence.

This weighting was introduced because the paper found that a fixed uncertainty-driven trajectory is not uniformly effective across tasks. On structured reasoning tasks such as GSM8K, the U-shaped pattern of confidence-based samplers is harmful because it tends to generate answers or formatting tokens early and backfill reasoning later. On 4×4 Sudoku, by contrast, a globally freer order can be beneficial because easy constraints can propagate inward from multiple positions (Huang et al., 18 Aug 2025).

The reported hyperparameter choices reflect this task dependence. For stepwise reasoning tasks, the paper uses $$\mathcal{L}_{\mathrm{MDM}} = - \mathbb{E}_{t,x_0,x_t}\left[ \frac{1}{t}\sum_{i=1}^L \mathbf{1}[x_t^i = \mathbf{M}] \log p_\theta(x_0^i \mid x_t) \right].$$3 for most tasks and $$\mathcal{L}_{\mathrm{MDM}} = - \mathbb{E}_{t,x_0,x_t}\left[ \frac{1}{t}\sum_{i=1}^L \mathbf{1}[x_t^i = \mathbf{M}] \log p_\theta(x_0^i \mid x_t) \right].$$4 for Countdown, which is described as benefiting from a more strongly left-to-right trajectory. For 4×4 Sudoku, the paper sets $$\mathcal{L}_{\mathrm{MDM}} = - \mathbb{E}_{t,x_0,x_t}\left[ \frac{1}{t}\sum_{i=1}^L \mathbf{1}[x_t^i = \mathbf{M}] \log p_\theta(x_0^i \mid x_t) \right].$$5, allowing the sampler to be position-neutral and rely only on informativeness (Huang et al., 18 Aug 2025). Heatmaps in the paper show that moderate positive $$\mathcal{L}_{\mathrm{MDM}} = - \mathbb{E}_{t,x_0,x_t}\left[ \frac{1}{t}\sum_{i=1}^L \mathbf{1}[x_t^i = \mathbf{M}] \log p_\theta(x_0^i \mid x_t) \right].$$6 converts the pathological U-shape into a near left-to-right pattern on reasoning tasks, while $$\mathcal{L}_{\mathrm{MDM}} = - \mathbb{E}_{t,x_0,x_t}\left[ \frac{1}{t}\sum_{i=1}^L \mathbf{1}[x_t^i = \mathbf{M}] \log p_\theta(x_0^i \mid x_t) \right].$$7 yields a more free-form trajectory on global planning tasks (Huang et al., 18 Aug 2025).

A plausible implication is that PC-Sampler treats position as a controllable inductive bias rather than an emergent property of local uncertainty. That distinguishes it from prior MDM decoding rules that rank masked positions solely by uncertainty and thereby inherit the model’s own preference for boundary tokens.

4. Confidence calibration and trivial-token suppression

The calibrated confidence term addresses the second limitation: trivial-token bias. In raw confidence-based decoding, a token can be selected early simply because it is globally frequent and therefore easy to predict. PC-Sampler discounts such predictions by multiplying model confidence with the negative log of the token’s empirical frequency in a background corpus (Huang et al., 18 Aug 2025).

The mechanism is frequency-sensitive in the following sense. If a token is highly frequent, then $$\mathcal{L}_{\mathrm{MDM}} = - \mathbb{E}_{t,x_0,x_t}\left[ \frac{1}{t}\sum_{i=1}^L \mathbf{1}[x_t^i = \mathbf{M}] \log p_\theta(x_0^i \mid x_t) \right].$$8 is large and $$\mathcal{L}_{\mathrm{MDM}} = - \mathbb{E}_{t,x_0,x_t}\left[ \frac{1}{t}\sum_{i=1}^L \mathbf{1}[x_t^i = \mathbf{M}] \log p_\theta(x_0^i \mid x_t) \right].$$9 is small, so even a large model confidence yields only a modest calibrated score. If a token is rare, then $>10\%$0 is larger, so a confident prediction of that token is amplified. The clipping threshold $>10\%$1 prevents extremely rare tokens from receiving unbounded scores; the paper uses $>10\%$2 across tasks and reports that performance is robust for moderate values (Huang et al., 18 Aug 2025).

The paper defines trivial tokens as high-frequency, semantically light tokens, including structural markers such as <|EOS|>, <|EOT|>, newline, and <SPACE>, punctuation, and common fillers such as “the”, “is”, and “answer is” (Huang et al., 18 Aug 2025). Under standard confidence-based decoding, the top-5 most frequent early-step selections are reported to be dominated by EOS, newline, and similar tokens, and more than $>10\%$3 of selected tokens in the first several steps come from the trivial set (Huang et al., 18 Aug 2025). PC-Sampler suppresses this behavior because high-frequency tokens receive small $>10\%$4 factors.

This calibration differs conceptually from entropy and margin scoring. Entropy and margin inspect the shape of the predictive distribution but remain oblivious to token informativeness. By contrast, PC-Sampler introduces an explicit corpus-level prior into the position score. This suggests a PMI-like bias toward tokens that are not merely predictable, but predictable relative to their baseline frequency. The paper notes that the design is inspired by divergence-based calibration used in pretraining-data detection, here repurposed as a token-level informativeness measure (Huang et al., 18 Aug 2025).

5. Experimental evaluation and qualitative behavior

The method is evaluated on three masked diffusion LLMs—LLaDA-Instruct-8B, LLaDA-1.5-8B, and Dream-v0-Instruct-7B—across seven benchmarks: GSM8K, MATH500, HumanEval, MBPP, GPQA, 4×4 Sudoku, and Countdown with 3 numbers (Huang et al., 18 Aug 2025). For comparison to autoregressive systems, the paper also reports results from LLaMA-3.1-8B-Instruct, Mistral-7B-Instruct, and Qwen-2.5-7B-Instruct (Huang et al., 18 Aug 2025).

Across LLaDA and LLaDA-1.5, PC-Sampler improves average performance by about $>10\%$5–$>10\%$6 absolute points over the best MDM baselines. On LLaDA-8B, the best baseline average is Fast-dLLM at $>10\%$7, while PC-Sampler reaches $>10\%$8. On LLaDA-1.5-8B, the best baseline average is Semi-AR at $>10\%$9, while PC-Sampler reaches $70\%$0. On Dream-7B, the best uncertainty-based baseline average is margin sampling at $70\%$1, while PC-Sampler reaches $70\%$2 (Huang et al., 18 Aug 2025).

Model	Best baseline average	PC-Sampler average
LLaDA-8B	37.0	42.3
LLaDA-1.5-8B	37.1	44.7
Dream-7B	27.8	40.1

Task-wise results also show gains. For LLaDA-1.5-8B, PC-Sampler reaches $70\%$3 on GSM8K versus a best baseline of $70\%$4, and $70\%$5 on MATH500 versus $70\%$6. For Dream-7B, it reaches $70\%$7 on HumanEval versus $70\%$8, and $70\%$9 on MBPP versus $t$0 (Huang et al., 18 Aug 2025). On the paper’s aggregate metric, LLaDA-1.5-8B with PC-Sampler slightly surpasses Qwen-2.5-7B-Instruct, with $t$1 versus $t$2, while maintaining the planning advantages of MDMs on tasks such as Sudoku and Countdown (Huang et al., 18 Aug 2025).

The ablation study decomposes the gain into trajectory control and confidence calibration. The paper compares a variant with only trajectory control, a variant with only calibrated confidence, and the full method. On GSM8K with LLaDA, the confidence baseline is $t$3, trajectory control alone reaches $t$4, confidence calibration alone reaches $t$5, and the full method performs best; on MBPP, the corresponding values are $t$6, $t$7, $t$8, and $t$9 (Huang et al., 18 Aug 2025). Both components are therefore complementary rather than redundant.

The qualitative case studies align with the quantitative pattern. On a GSM8K example, the baseline confidence sampler generates the final numeric answer early and then backfills reasoning, whereas PC-Sampler first decodes the intermediate reasoning steps and only later emits the final answer (Huang et al., 18 Aug 2025). On HumanEval, the baseline emits <EOS> and other trivial tokens early, producing incomplete code, whereas PC-Sampler delays EOS and allocates more tokens to function signatures, control flow, and the core algorithm (Huang et al., 18 Aug 2025).

6. Related formulations, limitations, and broader significance

PC-Sampler sits within a broader family of methods that combine position-sensitive control with confidence-sensitive decision rules, although those methods target different modalities and objectives. In streaming zero-shot TTS, “Chatterbox-Flash” introduces prior-calibrated scoring

$i \in \mathcal{M}_t$0

and an early-decoding schedule based on quantile thresholds over these calibrated scores (Seo et al., 29 May 2026). This suggests a close methodological parallel: both approaches correct raw position-wise confidence by subtracting or discounting a global prior before deciding which positions to commit.

In neural machine translation, “Confidence-Aware Scheduled Sampling” uses per-token confidence

$i \in \mathcal{M}_t$1

to determine whether each position should receive a ground-truth token, a predicted token, or a random token during the second decoding pass (Liu et al., 2021). The resulting policy is position-wise and competence-aware, but it operates in teacher-forced autoregressive training rather than inference-time masked diffusion decoding.

In ranking, “EviRank” extracts semantic, attention, and output evidences from a single forward pass, aggregates them through reliable opinion aggregation, and then applies a position-aware calibration

$i \in \mathcal{M}_t$2

to align confidence with ranking importance (Yan et al., 3 Jun 2026). A plausible implication is that PC-Sampler and EviRank exemplify the same design principle in different settings: local scores become more useful when they are explicitly calibrated against a notion of positional utility.

In noisy-label learning, “Enhanced Sample Selection with Confidence Tracking” replaces instantaneous confidence thresholds with trend statistics over per-sample confidence gaps, using the Mann–Kendall test to detect upward trajectories for correctly labeled but hard-to-learn samples (Pan et al., 24 Apr 2025). This suggests a broader interpretation of “confidence calibration” in which the calibration signal may be corpus frequency, ranking position, or temporal consistency, depending on the task.

The PC-Sampler paper also states several limitations. The optimal $i \in \mathcal{M}_t$3 is task-specific and currently chosen manually; the frequency calibration depends on a particular background corpus and may be suboptimal under domain shift; the trivial-token notion is still mostly frequency-based; the method does not provide explicit structural planning beyond a static positional prior; and experiments are confined to text benchmarks, leaving applications to other discrete domains open (Huang et al., 18 Aug 2025). Within those limits, the method is significant because it shows that decoding can materially narrow the gap between masked diffusion and autoregressive sequence models without changing the model architecture or training procedure (Huang et al., 18 Aug 2025).