Systems and methods for program synthesis
Embodiments described herein provide a reinforcement learning based framework engaging pretrained language models (LMs) for program synthesis tasks. Specifically, the framework adopts a training strategy that optimizes pretrained LMs for program synthesis tasks in an actor-critic approach.
1. A method of a reinforcement learning framework comprising a language model, an actor model and a critic model for code generation, the method comprising:
receiving, via an input interface, a problem specification in a natural language including unit tests which contain example input-output testing pairs corresponding to the problem specification and a corresponding code segment that solves a problem described in the problem specification;
finetuning the language model that has been pretrained on language tasks to generate code programs based on the problem specification and using the corresponding code segment as a ground truth;
generating, by the finetuned pretrained language model acting as the actor model, one or more tokens representing a sampled code program in response to the problem specification at a decoding time step;
compiling, by a processor, the sampled code program comprising the one or more tokens at a compiling environment and executing the sampled code program with a testing input from the unit tests to generate an execution result;
generating, by the critic model, a return indicating functional correctness of the sampled code program based on comparing the execution result of the sampled code program and a testing result of the problem specification;
computing a policy gradient of an expectation of the return given current parameters of the finetuned pretrained language model;
updating the current parameters of the finetuned pretrained language model according to the policy gradient; and
generating, by the finetuned pretrained language model having the updated current parameters, a repaired code program that repairs the sampled code program based on an input combining the problem specification, the sampled code program and the execution result of the sampled code program.
2. The method of claim 1 , wherein the sampled code program is generated by:
generating a predicted token for the sampled code program governed by the current parameters of the finetuned pretrained language model at the decoding time step; updating hidden state representations of the finetuned pretrained language model; and generating a next predicted token for the sampled code program using the updated hidden state representations of the finetuned pretrained language model at a next decoding time step.
3. The method of claim 2 , wherein the return is generated when an ending token is generated for the sampled code program.
4. The method of claim 1 , wherein the return is generated by: passing the sampled code program and a test to a complier; and determining a first return value depending on whether the sampled code program is successfully compiled and executed and whether the execution result matches with the testing result of the problem specification.
5. The method of claim 4 , wherein the policy gradient is computed as an estimate based on the return and a gradient of a conditional probability of a predicted token conditioned on prior predicted tokens and the problem specification.
6. The method of claim 4 , further comprising:
inputting, to the critic model, a baseline program generated by a base model in response to the problem specification; and
determining a second return value depending on whether the baseline program is successfully compiled and executed and whether the execution result matches with the testing result of the problem specification.
7. The method of claim 6 , wherein the wherein the policy gradient is computed based on a difference between the first return value and the second return value and a gradient of a conditional probability of a predicted token conditioned on prior predicted tokens and the problem specification.
8. The method of claim 1 , wherein the critic model is trained by: receiving, by the critic model, a training sequence of the problem specification and the sampled code program; generating, by the critic model, a predicted test outcome corresponding to the sampled code program; computing a cross-entropy loss by comparing the predicted test outcome and the execution result of the sampled code program; and updating the critic model based on the cross-entropy loss.
9. The method of claim 8 , wherein the predicted test outcome is computed a softmax operation of max-pooled contextual hidden states of a decoder in the critic model.
10. The method of claim 9 , wherein the policy gradient is computed based on a probability distribution of the predicted test outcome generated by the critic model and a gradient of a conditional probability of a predicted token conditioned on prior predicted tokens and the problem specification.
11. A system of a reinforcement learning framework comprising a language model, an actor model and a critic model for code generation, the system comprising:
an input interface that receives a problem specification in a natural language including unit tests which contain example input-output testing pairs corresponding to the problem specification and a corresponding code segment that solves a problem described in the problem specification;
a memory storing a plurality of processor-executable instructions;
a processor that reads and executes the plurality of processor-executable instructions to perform operations comprising:
finetuning the language model that has been pretrained on language tasks to generate code programs based on the problem specification and using the corresponding code segment as a ground truth;
generating, by the finetuned pretrained language model acting as the actor model, one or more tokens representing a sampled code program in response to the problem specification at a decoding time step;
compiling, by the processor, the sampled code program comprising the one or more tokens at a compiling environment and executing the sampled code program with a testing input from the unit tests to generate an execution result;
generating, by the critic model, a return indicating functional correctness of the sampled code program based on comparing the execution result of the sampled code program and a testing result of the problem specification;
computing a policy gradient of an expectation of the return given current parameters of the finetuned pretrained language model;
updating the current parameters of the finetuned pretrained language model according to the policy gradient; and
generating, by the finetuned pretrained language model having the updated current parameters, a repaired code program that repairs the sampled code program based on an input combining the problem specification, the sampled code program and the execution result of the sampled code program.
12. The system of claim 11 , wherein the sampled code program is generated by: generating a predicted token for the sampled code program governed by the current parameters of the finetuned pretrained language model at the decoding time step; updating hidden state representations of the finetuned pretrained language model; and generating a next predicted token for the sampled code program using the updated hidden state representations of the finetuned pretrained language model at a next decoding time step.
13. The system of claim 12 , wherein the return is generated when an ending token is generated for the sampled code program.
14. The system of claim 11 , wherein the return is generated by: passing the sampled code program and a test to a complier; and determining a first return value depending on whether the sampled code program is successfully compiled and executed and whether the execution result matches with the testing result of the problem specification.
15. The system of claim 14 , wherein the policy gradient is computed as an estimate based on the return and a gradient of a conditional probability of a predicted token conditioned on prior predicted tokens and the problem specification.
16. The system of claim 14 , wherein the operations further comprise:
inputting, to the critic model, a baseline program generated by a base model in response to the problem specification; and
determining a second return value depending on whether the baseline program is successfully compiled and executed and whether the execution result matches with the testing result of the problem specification.
17. The system of claim 16 , wherein the wherein the policy gradient is computed based on a difference between the first return value and the second return value and a gradient of a conditional probability of a predicted token conditioned on prior predicted tokens and the problem specification.
18. The system of claim 11 , wherein the critic model is trained by: receiving, by the critic model, a training sequence of the problem specification and the sampled code program; generating, by the critic model, a predicted test outcome corresponding to the sampled code program; computing a cross-entropy loss by comparing the predicted test outcome and the execution result of the sampled code program; and updating the critic model based on the cross-entropy loss.
19. The system of claim 18 , wherein the predicted test outcome is computed a softmax operation of max-pooled contextual hidden states of a decoder in the critic model, and
wherein the policy gradient is computed based on a probability distribution of the predicted test outcome generated by the critic model and a gradient of a conditional probability of a predicted token conditioned on prior predicted tokens and the problem specification.
20. A non-transitory machine-readable medium storing a plurality of processor-executable instructions for a reinforcement learning framework comprising a language model, an actor model and a critic model for code generation, the processor-executable instructions being executed by one or more processors to perform operations comprising:
receiving, via an input interface, a problem specification in a natural language including unit tests which contain example input-output testing pairs corresponding to the problem specification and a corresponding code segment that solves a problem described in the problem specification;
finetuning the language model that has been pretrained on language tasks to generate code programs based on the problem specification and using the corresponding code segment as a ground truth;
generating, by the finetuned pretrained language model acting as the actor model, one or more tokens representing a sampled code program in response to the problem specification at a decoding time step;
compiling, by the one or more processors, the sampled code program comprising the one or more tokens at a compiling environment and executing the sampled code program with a testing input from the unit tests to generate an execution result;
generating, by the critic model, a return indicating functional correctness of the sampled code program based on comparing the execution result of the sampled code program and a testing result of the problem specification;
computing a policy gradient of an expectation of the return given current parameters of the finetuned pretrained language model;
updating the current parameters of the finetuned pretrained language model according to the policy gradient; and
generating, by the finetuned pretrained language model having the updated current parameters, a repaired code program that repairs the sampled code program based on an input combining the problem specification, the sampled code program and the execution result of the sampled code program.