IC-1308Dynamic retrieval-based few-shot prompting substantially improves released LLMs' code optimization, with GPT-4-0613 reaching 76.07% optimization rate and 3.93x speedup (best@8), exceeding the 3.66x human reference
Alexander G Shypula, Aman Madaan, Yimeng Zeng, Uri Alon, Jacob R. Gardner, Yiming Yang, Milad Hashemi, Graham Neubig, Parthasarathy Ranganathan, Osbert Bastani, Amir Yazdanbakhsh
Using CodeBERTScore embeddings and FAISS to retrieve the k=4 most similar slow-fast pairs from the PIE training set, the paper constructs dynamic few-shot prompts for each test program. This strategy outperforms all baseline prompting variants across every model tested. GPT-4-0613 achieves 76.07% optimization rate and 3.93x speedup at best@8, surpassing the 3.66x human reference. CodeLlama 34B improves from 19.63% opt / 1.30x (CoT) to 42.54% opt / 2.43x with retrieval. The ablation over k in {1, 2, 4} shows k=4 is optimal for most models.
Increased speedup comes with some cost in correctness (e.g., CodeLlama 34B drops from 78.73% to 73.62% correct); retrieval uses the PIE training set, so the improvement is partly due to task-specific data rather than a general property of the model.