DDPT: Diffusion-Driven Prompt Tuning for Large Language Model Code Generation

TL;DR

The paper addresses the challenge of prompt engineering for NL2Code by proposing Diffusion-Driven Prompt Tuning (DDPT), which learns to generate optimal prompt embeddings from Gaussian noise while keeping the large language model frozen. A diffusion model is trained using the LLM's code generation loss to produce a directional vector that shifts the prompt embedding toward an optimal distribution, enabling rapid sampling from noise to effective prompts. Experiments on CodeAlpaca and CoNaLa with CodeT5p-based models show that DDPT often outperforms manual prompts and basic prompt-tuning across several metrics, and the embedding directions can be interpreted via nearest-neighbor words. The work highlights diffusion-based optimization as a promising approach for text-domain prompt engineering in NL2Code and discusses limitations and future directions, including variable-length prompts and controllable sampling.

Abstract

Large Language Models (LLMs) have demonstrated remarkable capabilities in code generation. However, the quality of the generated code is heavily dependent on the structure and composition of the prompts used. Crafting high-quality prompts is a challenging task that requires significant knowledge and skills of prompt engineering. To advance the automation support for the prompt engineering for LLM-based code generation, we propose a novel solution Diffusion-Driven Prompt Tuning (DDPT) that learns how to generate optimal prompt embedding from Gaussian Noise to automate the prompt engineering for code generation. We evaluate the feasibility of diffusion-based optimization and abstract the optimal prompt embedding as a directional vector toward the optimal embedding. We use the code generation loss given by the LLMs to help the diffusion model capture the distribution of optimal prompt embedding during training. The trained diffusion model can build a path from the noise distribution to the optimal distribution at the sampling phrase, the evaluation result demonstrates that DDPT helps improve the prompt optimization for code generation.

Figures (6)

• Figure 1: A diagram showing our diffusion-based training process. We first split each input sentence into context and instructions and transform them into word vector embeddings before feeding them into the diffusion model. Gray-colored components indicate Gaussian noise elements and the frozen LLM. GT stands for Ground Truth and text elements that are discrete are highlighted in pink.
• Figure 2: An illustration of our diffusion sampling procedure. The darker components represent the noisy vector embeddings, while the lighter components correspond to the progressively denoised embeddings. The dashed line indicates the modeled transition from the Gaussian noise distribution to the target optimal distribution.
• Figure 3: Sample code challenges from the CoNaLa dataset (blue) and CodeAlpaca20k dataset (purple). Though both datasets encompass diverse code generation tasks, they differ in expected output length: CoNaLa typically requires shorter code snippets, while CodeAlpaca20k prompts generally demand longer, more complex code solutions.
• Figure 4: The prompt template used in this study. We uses the default prompt template presented by the codeAlpaca dataset as our initial prompt representation. The template consists of three components: the context portion (orange), the user instruction (gray), and the generated code snippet (blue).
• Figure 5: Visual comparison of code generation results across different embedding methods. Orange denotes the original context embedding, green shows the optimized prompt embedding with its generated output, grey indicates code produced using original context embedding, and blue marks the ground truth code snippet.