Multi-Target Few-Shot Prompting with Diffusion Language Models

Anton Xue; Alan Li; Rahul Saha

cs.CLOct 15, 2026

Multi-Target Few-Shot Prompting with Diffusion Language Models

Authors: Anton Xue, Alan Li, Rahul Saha

Organizations: University of Texas at Austin

Abstract

Few-shot prompting is useful when instructions and rewards are hard to specify. One such application is editing by example, where a demonstrated formatting change or API migration may need to be applied across multiple sites. However, such edits often depend on both left and right context, making generation with causal autoregressive models awkward. We therefore introduce multi-target few-shot prompting with diffusion language models, where examples guide the infilling of multiple masked regions within a shared canvas. To address the known challenge of reliably following demonstrations in few-shot prompting, we further leverage selective masking of the canvas to derive a likelihood-ratio reward that favors example-supported infills. Across few-shot prompting benchmarks and software refactoring case studies, we find that demonstrations improve infilling accuracy, with example guidance yielding further gains in several settings.

Figures & tables

Figure 1: Multi-target few-shot prompting with a DLM. Demonstrations map event names and dates to initials, days, and abbreviated months. The DLM progressively fills two masked outputs within a shared canvas; tokens revealed at one site can inform subsequent predictions at the other.

Figure 2: Further examples of multi-target few-shot prompting. Two input–output demonstrations illustrate alphabetical sorting of word lists, which a DLM applies to two missing outputs.

Figure 3: Example guidance strengthens a demonstrated naming convention. With the demonstration output masked (

P_{0}

), the model favors item_count ; revealing it (

P_{1}

) increases support for itemCount , but does not change the ranking. At an illustrative

\beta=1

, example guidance (

Q_{\beta}

) amplifies this shift and favors itemCount . Probabilities are schematic and normalized over the two completions of the blue mask; [M] denotes a masked span.

Figure 4: Inference with unknown target lengths. A length predictor runs once on the compact canvas; each sample then draws a length allocation and generates a DLM infill over the expanded masks. Samples receive reward-guidance weights

w_{k}=(P_{1}(L^{(k)},u^{(k)})/P_{0}(L^{(k)},u^{(k)}))^{\beta}

, comparing their likelihoods with demonstration outputs visible (

P_{1}

) and masked (

P_{0}

). Resampling according to these weights approximates

Q_{\beta}

. Bars aggregate normalized weights over identical completions; tokens, samples, and weights are schematic.

Figure 5: Single-target prompting as the number of demonstrations grows. The top row shows unguided answer accuracy (

\beta=0

), while the bottom row shows the accuracy gain from example guidance (

\beta=0.5

), obtained by reweighting the same

K=64

samples per canvas. Four benchmarks and four model interfaces, with 64 prompt families per benchmark; DLMs use oracle target lengths, while Qwen3 generates freely.

Figure 6: Demonstrations and target count on List Functions–Easy. Each column fixes the number of jointly infilled targets

J

, while the horizontal axis varies the number of demonstrations

D

. The rows show all-target accuracy

A_{\mathrm{all}}

and average per-target accuracy

A_{\mathrm{avg}}

under unguided sampling (

\beta=0

), followed by the gain in

A_{\mathrm{all}}

from example guidance (

\beta=0.5

) using the same samples. Three DLMs, oracle target lengths, and

K=16

samples per canvas.

Figure 7: Length allocation and downstream infilling accuracy. On List Functions–Easy (

D=4

\beta=0

), the top row compares per-target accuracy

A_{\mathrm{avg}}

under oracle and sampled learned lengths as the number of targets

J

increases. The bottom row measures sensitivity to length errors, setting each target length to

L_{j}=L_{j}^{\star}+\Delta

for

J\in\{1,2\}

, where

L_{j}^{\star}

is the reference length and

\Delta=0

recovers oracle allocation. These forced offsets are distinct from the predictor’s own error distribution. Results cover three DLMs; additional length-prediction diagnostics appear in Section B.4 .

Figure 8: Java helper refactoring with LLaDA-1.5 (

K=16

). Left: mass of the Java length predictor within two tokens of the reference length versus a training-label prior, by reference length. Right: joint LLM-judged success by

(D,J)

, oracle or learned lengths, and unguided (

\beta=0

, hollow) or example-guided (

\beta=0.5

, filled) evaluation of the same sample pools. Pooled contrasts are in Table 1 .

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 9: Reference-token NLL under

P_{1}

and

P_{0}

by reveal position (cross-benchmark study,

J=8

). Previously revealed reference targets are visible in both contexts, and the gap between their subsequent reference-token scores generally narrows. This diagnostic follows reference assignments, not sampled decoding trajectories.

Figure 10: Cross-benchmark target-count study at the common guidance strength. Top: unguided all-target accuracy

A_{\mathrm{all}}

; bottom: its change under example guidance at

\beta=0.5

, which varies by task at

J=8

D=4

K=64

samples per canvas, oracle lengths, three DLMs, and 64 prompt families per cell; bands are unadjusted 95% family-bootstrap intervals.

Figure 11: Guidance-strength sensitivity on List Functions–Easy (cross-benchmark study). Change in

A_{\mathrm{all}}

from

\beta=0

J=1

and

J=4

(left), and mean

\rho_{\mathrm{eff}}

, the entropy-based effective fraction of sampled draws (right). Each strength reweights the same

D=4

K=64

sample pools, with duplicates retained and oracle lengths; Nemotron uses a causal prefix and bidirectional tail even at

J=1

. Bands are unadjusted 95% paired family-bootstrap intervals from 10,000 resamples.

Figure 12: Likelihood sharpening with the same total

P_{1}

exponent in the

D\times J

grid. Top: gain in

A_{\mathrm{all}}

from sharpening alone (sample weights

\propto P_{1}^{1/2}

, resulting distribution

\propto P_{1}^{3/2}

); bottom: example guidance minus sharpening (negative values favor sharpening). Both reweight the

K=16

sample pools of Figure 6 ; bands are unadjusted 95% family-bootstrap intervals.

Panel of Figure 13 — Figure 13: One editing canvas from a Checkstyle refactoring. An accepted edit introduces a helper call, and two masked targets must adapt it to their local arguments. Both columns belong to the same prompt; hook arrows indicate visual line wraps.

$(D,J)$	Commits/tasks	$\beta=0$	$\beta=0.5$	Difference	$\beta=0$	$\beta=0.5$	Difference	at $\beta=0$
		Oracle lengths			Learned lengths			Learned $-$ oracle
(1,1)	50/52	76.1	72.9	$-$ 3.2 [ $-$ 5.9, $-$ 0.8]	17.3	17.1	$-$ 0.2 [ $-$ 3.4, +3.9]	$-$ 58.9 [ $-$ 68.1, $-$ 49.5]
(1,2)	26/28	74.9	72.7	$-$ 2.2 [ $-$ 4.7, $-$ 0.2]	9.4	4.7	$-$ 4.6 [ $-$ 8.9, $-$ 1.3]	$-$ 65.5 [ $-$ 79.8, $-$ 50.5]
(2,1)	26/28	92.7	90.4	$-$ 2.2 [ $-$ 6.9, +0.9]	21.2	18.6	$-$ 2.5 [ $-$ 9.3, +5.9]	$-$ 71.5 [ $-$ 82.6, $-$ 59.1]
(2,2)	17/19	82.7	84.5	+1.8 [0.0, +5.2]	9.2	9.6	+0.4 [ $-$ 5.2, +8.0]	$-$ 73.5 [ $-$ 88.4, $-$ 56.6]
Pooled	50/127	74.4	71.5	$-$ 2.9 [ $-$ 5.4, $-$ 0.6]	15.8	14.5	$-$ 1.3 [ $-$ 3.8, +1.6]	$-$ 58.6 [ $-$ 67.4, $-$ 49.6]

Appendix

Table 1: Joint LLM-judged success (%) by configuration, with

K=16

samples per editing task and length setting and learned lengths sampled from the Java predictor. Differences are in percentage points, with unadjusted 95% paired commit-bootstrap intervals. The pooled row weights 50 commits equally and is not the mean of the four configuration rows.

Few-Shot Prompting

Few-shot prompting guides AI models with a small set of examples in the prompt, without updating their weights.

See all 43 papers

Unlocking Prompt Infilling Capability for Diffusion Language Models
Apr 4, 2026Yoshinari Fujinuma, Keisuke SakaguchiDiffusion Language Models Masked Diffusion Models
Patronus AI · Tohoku University, RIKEN
Match the Canvas, Not Just the Answer: Predictability-Aware SFT for Diffusion Language Models
Oct 13, 2026Miao Li, Hanyang Jiang, Hengyu Fu +8Diffusion Language Models Diffusion Model Fine-Tuning
Georgia Institute of Technology · University of California, Berkeley · University of Chicago +1
How Much Prompt Is Enough? A Blackbox Minimization of Few-Shots in LLMs
Sep 28, 2026Ali Alfageeh, Rahul Gopinath, Amin AlipourLLM Prompting Few-Shot Prompting
University of Houston · The University of Sydney

Apr 4, 2026J/K move · Enter open · S save

Unlocking Prompt Infilling Capability for Diffusion Language Models

Yoshinari Fujinuma, Keisuke Sakaguchi

Patronus AI · Tohoku University, RIKEN

Diffusion Language Models Masked Diffusion Models Diffusion Language Model Training Diffusion Model Fine-Tuning Prompt Engineering Prompt Optimization

Open paper AskPDF arXiv

Masked diffusion language models (dLMs) generate text through bidirectional denoising, yet this capability remains locked for infilling prompts. This limitation is an artifact of the current supervised finetuning (SFT) convention of applying response-only masking. To unlock this capability, we extend full-sequence masking during SFT, where both prompts and responses are masked jointly. Once unlocked, the model infills masked portions of a prompt template conditioned on few-shot examples. We show that such model-infilled prompts match or surpass manually designed templates, transfer effectively across models, and are complementary to existing prompt optimization methods. Our results suggest that training practices, not architectural limitations, are the primary bottleneck preventing masked diffusion language models from infilling effective prompts

Abstract

Few-Shot Prompting

Explore similar work

Unlocking Prompt Infilling Capability for Diffusion Language Models

Match the Canvas, Not Just the Answer: Predictability-Aware SFT for Diffusion Language Models

How Much Prompt Is Enough? A Blackbox Minimization of Few-Shots in LLMs