Diffusion language models improve few-shot multi-part text editing

Multi-Target Few-Shot Prompting with Diffusion Language Models

Computation and Language

Summary

Sometimes, teaching a computer to edit text by showing a few examples is tricky, especially when changes depend on the words before and after the edits. The authors found a way to use a special kind of language model called a diffusion model that can fill in multiple missing parts of text at once, guided by examples. They also designed a technique to better judge how well the model follows the examples, making the edits more accurate. This method works well for tasks like formatting text or updating code across many places.

What this means in practice

  • For software developers: Automatically update multiple related parts of software code based on a few example edits, improving correctness in API migrations.
  • For technical writers: Apply consistent formatting changes across large documents by showing a few examples of the desired output.

Authors

Anton Xue, Alan Li, Rahul Saha

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.