PAC Prediction Sets for Large Language Models of Code

TL;DR

This paper introduces a novel method for generating PAC prediction sets for large language models of code, providing high-confidence partial program outputs with theoretical guarantees, addressing the challenge of exponential label spaces in structured prediction.

Contribution

It proposes a new algorithm leveraging abstract syntax trees to produce compact, high-confidence partial programs for code generation models, a first in this domain.

Findings

Generates compact PAC prediction sets for code models

Demonstrates effectiveness on SQL and multiple programming languages

Provides theoretical guarantees for partial program predictions

Abstract

Prediction sets have recently been shown to be a promising strategy for quantifying the uncertainty of deep neural networks in a way that provides theoretical guarantees. However, existing techniques have largely targeted settings where the space of labels is simple, so prediction sets can be arbitrary subsets of labels. For structured prediction problems where the space of labels is exponential in size, even prediction sets containing a small fraction of all labels can be exponentially large. In the context of code generation, we propose a solution that considers a restricted set of prediction sets that can compactly be represented as partial programs, which are programs with portions replaced with holes. Given a trained code generation model, our algorithm leverages a programming language's abstract syntax tree to generate a set of programs such that the correct program is in the set with high-confidence. Valuable applications of our algorithm include a Codex-style code generator with holes in uncertain parts of the generated code, which provides a partial program with theoretical guarantees. We evaluate our approach on PICARD (a T5 model for SQL semantic parsing) and Codex (a GPT model for over a dozen programming languages, including Python), demonstrating that our approach generates compact PAC prediction sets. This is the first research contribution that generates PAC prediction sets for generative code models.