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 ( P0 ), the model favors item_count ; revealing it ( P1 ) increases support for itemCount , but does not change the ranking. At an illustrative β=1 , example guidance ( Qβ ) 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 wk=(P1(L(k),u(k))/P0(L(k),u(k)))β , comparing their likelihoods with demonstration outputs visible ( P1 ) and masked ( P0 ). Resampling according to these weights approximates Qβ . 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 ( β=0 ), while the bottom row shows the accuracy gain from example guidance ( β=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 Aall and average per-target accuracy Aavg under unguided sampling ( β=0 ), followed by the gain in Aall from example guidance ( β=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 , β=0 ), the top row compares per-target accuracy Aavg 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 Lj=Lj⋆+Δ for J∈{1,2} , where Lj⋆ is the reference length and Δ=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 ( β=0 , hollow) or example-guided ( β=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 P1 and P0 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 Aall ; bottom: its change under example guidance at β=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 Aall from β=0 at J=1 and J=4 (left), and mean ρ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 P1 exponent in the D×J grid. Top: gain in Aall from sharpening alone (sample weights ∝P11/2 , resulting distribution ∝P13/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.
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.
Oracle lengths
Learned lengths
Learned − oracle
(D,J)
Commits/tasks
β=0
β=0.5
Difference
β=0
β=0.5
Difference
at β=0
(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.
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
Diffusion language models (DLMs) enable flexible parallel generation by resolving tokens according to the model's evolving predictability rather than a fixed left-to-right order. This flexibility, however, introduces training choices absent from standard autoregressive teacher forcing: which reference tokens are visible when a target is supervised and how much generation space is allocated. Standard supervised fine-tuning (SFT) typically makes these choices through random masking, independently of inference. Our analysis shows how this training-inference mismatch can misdirect supervision. Random masking can conceal prediction errors or penalize uncertainty that decoding would otherwise resolve, while the attention-visible input length can change which response formulations the model favors, even when the underlying solution is unchanged. Motivated by these findings, we propose Predictability-Aware Canvas Selection and Supervision (PACS), an SFT method. PACS uses the model's own predictability to select each example's attention-visible input length and supervises reference tokens immediately before their confidence-guided revelation. Across reasoning and code-generation benchmarks, PACS consistently achieves the highest mean accuracy among the compared SFT methods, with gains of up to 12.84 percentage points over vanilla SFT on individual benchmarks. Code will be released upon publication.
Miao Li, Hanyang Jiang, Hengyu Fu +8
Georgia Institute of Technology · University of California, Berkeley · University of Chicago +1
Prompts are the primary mechanism for directing the behavior of large language models (LLMs). Yet the internal structure and causal hierarchy of prompts remain poorly understood: which parts are causally necessary and which are redundant is an open question. This opacity can have severe consequences. Subtle prompt variations can silently shift model outputs in critical software systems, and engineers lack techniques to reason about prompt reliability. We present \framework, a blackbox prompt-minimization framework that reduces few-shot prompts to their necessary minimal subset. We use a case study to apply \framework to a few-shot learning system and demonstrate the insights that this framework can provide. Our experiments show that few-shot exemplars can be reduced by a mean of 65.3%~±~15.8% in character count while fully preserving propositional output fidelity. The models preferentially retain logical identifiers and constraint declarations while discarding natural language prose and cross-prompt relational annotations. Our analysis also shows that some models are universal encoders, able to produce highly legible yet minimized prompts, while others are universal decoders, able to interpret minimized prompts from most other models. By identifying which components are indispensable, \framework provides a principled basis for prompt compression and structural analysis of few-shot exemplars.