\section{Conclusion}

Although large language models have shown promise towards auto-coding,
the low success rate of existing language models remains an open problem. In this paper, we study the mistakes made by auto-generated programs by Codex, and  investigate whether existing APR tools and the newly released Codex-e can fix the auto-generated buggy programs.
Our experiment results reveal that: (1) program produced by Codex share common defect categories as human programmers; (2) existing APR techniques (TBar and Recoder) do not perform well at fixing bugs in auto-generated programs

(3)  given proper instructions (information from fault localization), Codex-e shows promising initial results in code edit generation , which outperforms TBar and Recoder by fixing 45\% more buggy solutions. The implications of our study include: (1) enhancing language models using traditional software engineering techniques (e.g., AST information, fault localization), (2) highlighting the limitations of APR techniques, especially pattern-based approaches, and (3) suggesting future direction of APR research (e.g., flexible form of fault localization). The ability to generate multi-hunk patches at flexible locations can be a milestone for future (learning-based) APR research.

\section{Introduction}
\label{sec:intro}

Designing AI-based systems to automatically solve programming tasks has gained considerable attention in recent years.
The most notable of these comes in the form of transformer-based large-scale language models, which used to achieve impressive performance in generating text.
The transformer-based models, such as Codex~\cite{codex} and AlphaCode~\cite{alphacode}, have successfully generated code for many programming tasks in Python, Java, and C.
Technically, those techniques treat code generation as a transformation problem, which takes as input natural language descriptions and transforms them into programming language.

Although transformer-based language models successfully solved many programming tasks, their success rate is still relatively low.
When evaluating on \emph{pass@5} metric~\cite{codex}, the best Codex model achieves ~24.52\% passing rate at introductory-level tasks and ~3.08\% passing rate at competition-level tasks~\cite{codex} from APPS dataset~\cite{apps}.
While the best AlphaCode model achieves 20.36\% and 7.75\% passing rates on introductory-level and competition-level tasks respectively~\cite{alphacode}.
Lacking deep understanding of task descriptions and program semantics are the main reasons that cause the low success rate.
Transformer-based models treat code generation as a sequence-to-sequence transformation by treating description and code as token sequences which cannot capture deep semantic features of programs (e.g., program semantic can be encoded in the form of a test case).
In contrast, generating entire programs requires understanding the entire task descriptions which usually comprise complex logic, and figuring out the solutions to programming tasks relies on deep algorithm reasoning.

\input{table/keyfindings}

Automated program repair (APR) is an emerging area for automated rectification of programming errors~\cite{apr}.
Automated repair techniques take as inputs a buggy program and a correctness specification, and produce a fixed program by slightly changing the program to make it satisfy the given specification.
Typical repair tools generate patches by reasoning about the program semantics against the given specification.
For instance, semantic-based repair tools (e.g., SemFix~\cite{semfix}, Angelix~\cite{angelix}) generate patches by using symbolic execution and search-based repair tools (e.g., GenProg~\cite{genprog}, TBar~\cite{tbar}) search for correct patches among pre-defined search space with the reference to dynamic execution results.
APR has shown promising results in fixing real-world bugs and vulnerabilities.
For instance, in a well-studied dataset Defects4J~\cite{defects4j} version 1.2.0 including 395 real bugs, more than 25\% (101) of bugs have been automatically fixed by at least one APR tool~\cite{tbar}.
Nevertheless, existing APR tools are limited to generating small patches (usually one-line fixes) due to the complexity of semantic reasoning (semantic-based APR) and search space explosions problem (search-based APR) when considering multiple-line fixes.
In other words, APR tools are good at generating small patches via semantic reasoning but cannot generate a large chunk of code.

The strength and weakness of language models and APR techniques inspire us to think about the following question: \textit{can automated program repair improves the code produced by language models?}

In this paper, we apply existing APR techniques to the code generated by the Codex model, and answer the following research questions:

\textit{\textbf{(RQ1) What mistakes do the auto-generated code usually make?}}

Although we know that the language models produce many wrong solutions when solving programming tasks, several open questions still remain: (i) what are the types of bugs made by language models; (ii) are the bugs made by language models similar with human-written bugs; (iii) how many of these bugs can be fixed with single-hunk change.
We first study the bug pattern of code produced by Codex on how they can be fixed.

\textit{\textbf{(RQ2) How effective are APR tools in fixing code from Codex?}}

Existing APR tools are mainly designed to fix human-written bugs.
To achieve this, common APR tools generate patches by defining transformation operators (search-based APR) or specifying the program synthesis ingredients (semantic-based APR).
These operators and ingredients have been proven to be efficient in fixing human-written bugs.
We study whether these operators or ingredients still work when fixing the bugs in auto-generated code.
Specifically, we study how effective APR tools are in fixing the code produced by Codex.

\textit{\textbf{(RQ3) Can Codex edit mode fix program bugs?}}

In March 2022,  a new version of Codex was released, which can edit existing content in a complete program rather than just completing a partial program~\footnote{https://openai.com/blog/gpt-3-edit-insert}.
This new feature makes it practical to use Codex to modify or improve existing code for program transformation, code reconstruction, and bug fixing.
Codex edit mode (we call this mode \textit{Codex-e} throughout this paper) requires users to provide instructions to guide the revision, such as ``translate the java program to javascript''.
To fix a bug, users need to provide precise and clear instructions.
How to automatically produce such instruction still remains an open question.
We study whether the side effect of APR tools, such as fault localization results, can be used to guide Codex-e, and how efficient Codex-e is in fixing program bugs.

The key findings of our study are summarised in Table~\ref{tab:findings}.
According to these key findings, we discuss how to further improve the success rate of language models.
In this paper, we mostly used pattern-based and learning-based APR tools in our experiments.
Our result shows that existing APR tools are still quite limited, including limited patch space, fix locations and patch size, enhancing APR tool to surpass these limitations (e.g., introducing a more flexible fault localization strategy) is highly desirable.
Besides those repair tools, there are many other semantic analysis techniques that are widely-used in traditional software engineering could be further investigated for improving the efficiency of language models.

Overall, our contributions can be summarized as follows.
\begin{itemize}[leftmargin=*, topsep=0pt]
\item To the best of our knowledge, we present the first study of automated repair of buggy programs automatically generated by the Codex model.
\item  To the best of our knowledge, we include the first study that evaluates the efficacy of the newly released Codex-e as an automated repair tool.
\item We propose LMDefects{}, a new dataset that contains 42 correct Java programs (37 from easy-level tasks and 5 from medium-level tasks) and 355 buggy Java programs (115 from easy-level tasks and 240 from medium-level tasks), together with a taxonomy of the defect types exist in these buggy solutions.
\end{itemize}

\section{Related Work}

In this section, we present existing APR techniques that are relevant to our study and the language models for code generation.

\noindent
\textbf{Automated Program Repair.}
Automated Program Repair (APR) has gained a lot of attentions from both academia and industry in recent years~\cite{cacm19}.
The widely-studied APR techniques include search-based, semantic-based and learning-based APR.

\textit{Search-based APR.}
Search-based APRs~\cite{relifix,droix, capgen, simfix, F1X, arja, tbar, liu2019avatar}, such as GenProg~\cite{genprog}, takes a buggy program and a correct criteria as inputs, and generates patches in two steps: (1) generating a candidate patch space using predefined code transformation operators; and (2) searching for the correct patch over the patch space that satisfies the given correctness criteria.
Search-based repair is able to fix many real-world bugs and can scale to large programs.
However, the patch space is limited by the pre-defined fix patterns. At the same time,
defining large number of fix patterns causes the search space explosion problem.
In this work, we studied the effectiveness of search-based APR in fixing auto-generated code.
\textit{Learning-based APR.} The application of deep learning techniques in program repair has been explored in past few years. DeepRepair~\cite{deeprepair} and DeepFix~\cite{gupta2017deepfix} are the early attempts to fix bugs by learning fixes from similar code. SequenceR~\cite{chen2019sequencer} adapts neural machine translation (NMT) to generate patch, where CoCoNuT~\cite{lutellier2020coconut} and CURE~\cite{jiang2021cure} further improve the results by either encoding program context or using a programming language model. DLFix~\cite{li2020dlfix} use two-layer tree-based RNN to learn code transformations, and Recoder~\cite{recoder} designed a syntax-guided learning approach to improve the decoder of a DL model. In this work, we select Recoder because it fixes the most number of bugs in Defects4J~\cite{defects4j} among those DL-based APR tools with the trained model released.

\textit{Semantic-based APR.}
Semantic-based APR techniques, such as
SemFix~\cite{semfix}, Nopol~\cite{NOPOL}, and Angelix~\cite{angelix}, generate patches by (1) formulating a repair constraint that needs to be satisfied by a program passing a given test-suite; and (2) solving the repair constraint to generate patches. In this work, we have only studied search-based and learning-based repair tools; semantic based APR tools have not been covered in our work.

\noindent \textbf{Large Language Model for Code Generation.} Large language models such as GPT-3~\cite{gpt3} have shown promising performance in the NLP domain. Hendrycks et al.~\cite{apps} proposed APPS dataset and evaluated the code generation performance of several variant GPT models with APPS as the fine-tuned data. Later, Codex~\cite{codex}, the back-end model that powers GitHub Copilot improved the results by fine-tuning GPT-3 with much larger training data in GitHub. AlphaCode~\cite{alphacode} is similar to Codex, but focus more on producing program solutions for difficult programming tasks. There are also emerging works combining program synthesis with large language model~\cite{jigsaw, googlellmsynthesis, rahmani2021multi}. The most relevant papers to us are studies on how language model can fix bugs~\cite{codex_security, codex_apr}. In contrast, we evaluated whether APR tools (including Codex-e) can fix programs automatically produced by Codex.

\section{RQ1: What mistakes do auto-generated code usually make?}
\label{sec:rq1}

\begin{table*}[!t]
\caption{Defect Classification of Unsolved Solutions}
\label{tab:bug-pattern}
\small
    \begin{tabular}{l|p{4.8cm}|l|r|r|r}\hline
        Defect Category & Sub-category & Example & Easy & Medium & Total\\\hline
\hline
\multirow{3}{*}{ \textbf{M}ulti-hunk}
        & (M-S) \textbf{S}imilar  &
\begin{lstlisting}[style=mystyle, numbers=none, frame=none]
-for (int n : nums1) {
+for (int n : set1) {
-for (int n : nums2) {
+for (int n : set2) {
\end{lstlisting}
 & 4 & 1 &  5\\\cline{2-6}

        & (M-U) \textbf{U}nique & \begin{lstlisting}[style=mystyle, numbers=none, frame=none]
-for (int i = 2; i < s.length(); i++)
+for (int i = 0; i < s.length(); i++)
-if (s.charAt(i) == 'X' && ExprA)
+if (s.charAt(i) == 'X')
\end{lstlisting}
& 4 & 8  & 12\\\cline{2-6}
        & (M-L) Need \textbf{L}arge Fix & Refer to Supplementary & 6 & 14 &  20\\\hline
\multirow{8}{*}{\textbf{S}ingle-Hunk}
        & (S-O) \textbf{O}perator Mutation  &
\begin{lstlisting}[style=mystyle, numbers=none, frame=none]
-if (original.length < m * n)
+if (original.length != m * n)
\end{lstlisting}
        &2 & 3  & 5  \\\cline{2-6}
        & (S-V) \textbf{V}ariable Mutation &
\begin{lstlisting}[style=mystyle, numbers=none, frame=none]
-count++;
+ans++;
\end{lstlisting}
        & 3 & -  & 3  \\\cline{2-6}

        & (S-T) \textbf{T}ype Cast &
\begin{lstlisting}[style=mystyle, numbers=none, frame=none]
-Set<String> quadruplets = new HashSet<>();
+List<String> quadruplets = new ArrayList<>();
\end{lstlisting}
        &1 & -  & 1   \\\cline{2-6}

        & (S-E)  \textbf{E}xpression Mutation on Operand &
\begin{lstlisting}[style=mystyle, numbers=none, frame=none]
-i-=2;
+i-=1;
\end{lstlisting}
        &7 & 3  & 10   \\\cline{2-6}

        &(S-AS) \textbf{A}dd \textbf{S}tatement &
\begin{lstlisting}[style=mystyle, numbers=none, frame=none]
+Arrays.sort(res);
\end{lstlisting}
        & 8 & 2  & 10   \\\cline{2-6}
        &(S-MS) \textbf{M}ove \textbf{S}tatement &
\begin{lstlisting}[style=mystyle, numbers=none, frame=none]
-for (int j = row1; j <= row2; j++) {
 for (int i = col1.charAt(0);...;...) {
+for (int j = row1; j <= row2; j++) {
\end{lstlisting}
        &1 & -  &  1  \\\cline{2-6}
        &(S-DS) \textbf{D}elete \textbf{S}tatement &
\begin{lstlisting}[style=mystyle, numbers=none, frame=none]
-i++;
\end{lstlisting}
        &1 & 1  &  2 \\\cline{2-6}
\hline
\multirow{2}{*}{Syntax Error}
& Incomplete Code  & Refer to Supplementary & 35 & 92 & 127 \\\cline{2-6}
& Bracket Mismatch & Refer to Supplementary & 10 & 19 & 29 \\\cline{2-6}
& Invoke undefined functions/classes & Refer to Supplementary & 7 & 4 & 11 \\\cline{2-6}
& Incompatible types  & Refer to Supplementary & - & 1 & 1 \\\hline
               Algorithm-related
        & Misaligned Algorithm &Refer to Figure~\ref{listing:example-dp}  &26 & 92 & 118{} \\\hline
\rowcolor{grey}
    Total   & - & - & 115 & 240 & 355\xspace\\
        \hline

\end{tabular}
\end{table*}

Prior study of search-based automated program repair techniques shows that automatically generated patches are likely to exhibit certain \emph{anti-patterns} (program transformations that should be prohibited as they lead to the generation of nonsensical patches)~\cite{tan2016anti}. Intuitively, automatically generated code by a large language model like Codex may also contain anti-patterns. Hence, we analyze whether the code generated by Codex made the same common mistakes. Specifically, when we give a programming task in LMDefects{} for Codex to solve, we first run the five automatically generated solutions on the public tests and then validate those solutions that pass public tests on the private tests by submitting them to LeetCode~ online judge platform. If all the five automatically generated solutions $S$ by Codex cannot pass all tests (public and private tests), we consider $S$ as \emph{unsolved solution}. If there is no \emph{unsolved solution} $S$ for a programming task, we consider the programming task as \emph{solved}. Finally, Codex produces 37\xspace and 5\xspace \emph{solved} programming tasks on easy and medium levels respectively, and we study the mistakes of 355\xspace \emph{unsolved solution} $S$ in the remaining 71\xspace \emph{unsolved} programming tasks that lead to compilation errors or test failures.

For each unsolved solution $S_{buggy}$ generated by Codex, we manually fix the bugs in the solution by first referring to other solutions for the same programming task for repair hints, and then constructing a minimal patch that fixes the bugs by making the least program modifications to existing code. Our goal is to construct a ``ground truth'' patch $S_{fixed}$ for each unsolved solution to obtain
the ``diff'' between $S_{buggy}$ and $S_{fixed}$. This ``diff'' represents the bugs or mistakes in the automatically generated program by Codex. Based on this ``diff'', we manually classify each unsolved solution using the following categories:

\begin{itemize}[leftmargin=0pt]
\item   \textbf{\emph{Multi-hunk fix required}}

   \begin{description}[leftmargin=0pt]
        \item[(M-S) Similar:] Similar bugs (require similar fixes) exist at multiple program locations
        \item[(M-U) Unique:] Distinct bugs exist at multiple program locations
         \item[(M-L) Need Large Fix:] Need to edit more than five lines in total at multiple program locations
    \end{description}
    \item \textbf{\emph{Single-hunk fix required}}
       \begin{description}[leftmargin=0pt]
        \item[(S-O) Operator Mutation:] Change the arithmetic/logical/bitwise operators
         \item[(S-V) Variable Mutation:] Replace with a different variable
        \item[(S-T) Type Cast:] Change the data type
        \item[(S-E) Expression Mutation on Operand:] Modify an expression as part of an operand
        \item[(S-AS) Add Statement:] Insert a new statement
        \item[(S-MS) Move Statement:] Move a statement to a new location
        \item[(S-DS) Delete Statement:] Delete an existing statement

        \end{description}
      \item \textbf{\emph{Syntax error}}
   \begin{description}[leftmargin=0pt]
        \item[Incomplete Code:] For the last line of the program, only parts of the program is printed
        \item[Bracket mismatch:] Fails to compile due to missing/extra bracket
        \item[Invoke undefined functions/classes:] Fails to compile due to invoking a function or a class that do not exist
        \item[Incompatible types:] Fails to compile due to type errors
     \end{description}
     \item \textbf{\emph{Algorithm-related}}
       \begin{description}[leftmargin=0pt]
        \item[Misaligned algorithm:] The algorithm used is misaligned with the requirement given in the task description.
        \end{description}
\end{itemize}

Table~\ref{tab:bug-pattern} shows the defect classification of unsolved solutions, together with examples (the ``Example'' column), and the difficult levels (the ``Easy'' and the ``Medium'' columns) of the programming tasks to explain each defect. To allow for easier comparison of defect types, we derive the defect classification based on categories used in Codeflaws~\cite{tan2017codeflaws} (a benchmark that contains incorrect submissions by participants in programming competitions). Compared to the defect classification in Codeflaws, we realized that the types of bugs made by automatically generated code overlap with those in Codeflaws. Specifically, both Codeflaws and our dataset contain defects where either multi-hunk or single-hunk fixes are required. Moreover, for the single-hunk fixes, both datasets share similar mutation operators (e.g., operator mutation, and variable mutation). This indicates \emph{Codex made similar programming mistakes as human participants for defects that require multi-hunk or single-hunk fixes}. We think this is expected because Codex is trained with a lot of human-written programs that can be potentially buggy.

Table~\ref{tab:bug-pattern} shows that most defects made by Codex are syntax errors and algorithm-related errors. In fact, when the difficulty levels for programming tasks increase from ``Easy'' to "Medium", the number of syntax errors and algorithm-related errors increase, indicating that \emph{Codex is less likely to solve a more difficult programming task compared to a problem with a lower difficulty level}. This observation is in line with prior study which shows that the AlphaCode models perform better on programming tasks with lower difficulty ratings than in more difficult problems~\cite{alphacode}.

Given that syntax errors and algorithm-related errors are prevalent in Codex generated solutions, we manually analyzed these solutions to study the root causes behind these errors.

\noindent \textbf{Syntax Errors.} Our manual analysis revealed that automatically generated programs that lead to compilation errors usually have (1) incomplete code, (2) missing a closing bracket or adding an extra closing bracket; (3) invoking an undefined functions or classes or (4) have incompatible types. Some of these compilation errors (e.g., bracket mismatch) are simple enough to be automatically fixed. For programs with undefined functions, we need to synthesize the function body to resolve the compilation errors. For automatically generated programs with incomplete code, the last line of code is only partially printed. This may be due to the fact that the Codex model is restricted by the maximum number of tokens, which leads to incomplete code being generated for longer programs. In fact, programs with incomplete code have 13-46 lines of code, indicating that Codex may have difficulty generating longer programs.

\noindent \textbf{Misaligned Algorithm.}  Among all unsolved solutions, 118{} solutions use misaligned algorithms to solve the given programming tasks. Figure~\ref{listing:example-dp} shows an example of an automatically generated solution that uses misaligned algorithm. Given the prompt at lines 1--8, the Codex model incorrectly solve the \texttt{minimumOperations} task using a well-known dynamic programming (dp) algorithm. We think that Codex fails to solve this problem because the correct solution requires the model to generate a new customized solution for the task, whereas a GitHub code search using the query ``minimumOperations'' suggests several solutions that use dynamic programming. This indicates that Codex may return an incorrect solution, which is misled by the method name (\texttt{minimumOperations}). The problem of generating solutions that do not meet the user intention is known as the \emph{misalignment} problem in the Codex paper~\cite{codex}. All defects classified as ``misaligned algorithm'' suffer from the misalignment problem.

\begin{tcolorbox}
\vspace{-3pt}
Auto-generated programs share common mistakes with human-written programs but there are specific prevalent mistakes (syntax errors and misaligned algorithm).
\vspace{-3pt}
\end{tcolorbox}
\section{RQ2: How effective are APR tools in fixing the code produced by Codex?}
\label{sec:rq2}

As shown in Section~\ref{sec:rq1}, the mistakes made by Codex share some similarities with human-written solutions.
In this section, we study how effective existing APR tools are in fixing the code produced by Codex.
Since existing APR tools are designed to produce small patches (usually one-line or few-line fixes), we exclude the buggy programs that cannot be compiled or require changing the entire algorithms, which are not supported by current APR tools.

We are aware of techniques that can automatically fix compilation errors~\cite{ahmed2018compilation, deepfix}.
However, in this study, we only evaluate tools that fix coding errors, and we leave the evaluation of tools for fixing compilation errors as future work.

Given the unsolved solutions by the Codex model (excluding solutions that produce syntax errors and algorithm-related errors), we run TBar and Recoder to assess their ability in generating patches. During the patch validation stage, the automatically generated patches are categorized as below:

\begin{table}[!t]
\caption{The number of generated patches and the number of fixed tasks by TBar and Recoder (include single-hunk and multi-hunk)}
\label{tab:correct_vs_plausible}
    \begin{tabular}{l|cc|cc}
    \hline
    Tool & \multicolumn{2}{c|}{Correct/Plausible patches} & \multicolumn{2}{c}{Correctly Fixed Tasks}\\\cline{2-5}
     & easy & medium & easy & medium \\\hline
    TBar & 4/12 & 2/8 & 3 & 2\\\hline
    Recoder & 6/10 & 4/12 & 4 & 4\\\hline

\end{tabular}
\end{table}

\begin{description}[leftmargin=*]
\item [Plausible patches.] Plausible patches are patches that make the unsolved solutions pass the given public tests.
\item[Correct patches.] Correct patches are patches that make the unsolved solutions pass both the public tests and private tests and accepted by LeetCode{}.

\end{description}

\begin{figure}[!t]
\centering
\begin{lstlisting}[style = mystyle, escapechar=@]
// task delete-characters-to-make-fancy-string from biweekly-contest-58
// The patch produced by Recoder
public static String makeFancyString(String s){...
    if(...) {
      sb.deleteCharAt(i);
-      i -= 2;
+      i -= 1; // expression mutation (S-E-6)
    }...
}

// task watering-plants from weekly-contest-268
// The patch produced by Recoder
public static int wateringPlants(int[] plants, int capacity) {...
  if (plants[i] > currWater) {
    steps += (i - 1) * 2;
+   steps++; //add a statement (S-AS-10)
    ...} ...
}
\end{lstlisting}
\vspace{-6pt}

\caption{Examples of LeetCode programming tasks fixed by Recoder but not TBar.}
\label{listing:tbar_vs_recoder}
\vspace{-6pt}
\end{figure}

Table~\ref{tab:correct_vs_plausible} shows the number of generated patches and the number of correctly fixed programming tasks by TBar and Recoder, respectively.
Although TBar produces 12 and 8 plausible patches on easy-level and medium-level tasks, it only correctly fixes 4 easy and 2 medium patches.
Compared to TBar, Recoder produces more plausible patches (10 and 12 on easy and medium level, respectively), and more correct patches (6 and 4).
If we compare the percentage of correct patches among all generated patches, Recoder still outperforms TBar (6/20*100=30\% for TBar versus 10/22*100=45.5\% for Recoder).
As shown in prior studies on APR tools~\cite{smith2015cure,yi2018correlation}, test cases play an important role in patch generation.
Hence, we think that the limited number of public tests is one of the reasons that prevent APR tool from generating more correct patches.

The ``Correctly Fixed Tasks'' columns of Table~\ref{tab:correct_vs_plausible} show the number of programming tasks correctly fixed by TBar and Recoder, respectively. Note that each programming task corresponds to the five selected unsolved solutions.
If any of solutions is correctly fixed (accepted by LeetCode{}), we consider that this task has been solved (after five trials). In total, Recoder fixes eight programming tasks whereas TBar only fixes five tasks.
Overall, TBar increases the number of solved tasks from 37\xspace to 40 on the easy-level tasks (Recoder further increase this number by fixing two other easy-level tasks), while Recoder increases the number of solved medium-level tasks from 5\xspace to 9.
Combining both tools, APR tools help Codex fix 5 and 4 more easy-level and medium-level tasks, respectively.

We further analyze the type of defects fixed by the two APR tools. Table~\ref{tab:data_apr} shows the number of solutions that can be correctly fixed for each defect category, where the ``TBar'' and the ``Recoder'' columns show the number of patches produced by the corresponding tools.
For each category, the repair tools may not fix the bug by minimally changing the program (i.e., repair tools may fix a bug using different operators than the minimal fix shown in the ``Defect sub-category'' column). Overall, on the easy-level tasks that require single-hunk fixes, TBar correctly fixes 4 out of 23 solutions, while on medium-level tasks, TBar fixes 2 out of 9 buggy solutions.
In contrast, Recoder correctly fixes 6 out of 23 easy solutions, while on medium-level tasks, it fixes 4 out of 9 buggy solutions. Figure~\ref{listing:tbar_vs_recoder} shows two examples where Recoder outperforms TBar. In the first example, despite having the ``Mutate Literal Expression'' pattern, TBar fails because it cannot find the correct literal to replace due to limited mutation space (e.g., only support change 1 to 1d or 1f, but we can use program synthesis technique to find the correct literal). For the second example, TBar fails to generate the correct patch because it does not have any patterns that insert statements by copying from other locations. For the tasks which require multi-hunk fixes, both TBar and Recoder fail to generate any correct patches. This shows that \emph{existing APR tools are still quite limited in generating complex patches that require edits of multiple lines}.

\begin{table}[!t]
\caption{The number of correctly fixed solutions using different APR tools (only single-hunk bugs considered).}
\setlength{\tabcolsep}{2pt}
    \begin{tabular}{l|cc|cc|cc}
         \hline
        Defect  & \multicolumn{2}{c|}{Total} & \multicolumn{2}{c|}{TBar}  & \multicolumn{2}{c}{Recoder} \\\cline{2-7}
        Sub-category& easy & medium & easy & medium & easy & medium                \\\hline\hline
        S-O   & 2 & 3  & 1 & 0  & 1 & 0 \\\hline
        S-V   & 3 & -  & 1 & -  & 1 & - \\\hline
        S-T   & 1 & -  & - & -  & - & - \\\hline
        S-E   & 7 & 3  & 1 & 2  & 2 & 2 \\\hline
        S-AS  & 8 & 2  & - & -  & 1 & 2 \\\hline
        S-MS  & 1 & -  & - & -  & - & - \\\hline
        S-DS  & 1 & 1  & 1 & -  & 1 & - \\\hline

\rowcolor{grey}
        Total  (Single-Hunk) & 23& 9  & 4 & 2  & 6 & 4  \\\hline

M-S/M-U/M-L   & 14    & 23      & -    & -      & -    & -           \\\hline
\end{tabular}
\label{tab:data_apr}
\end{table}

\begin{tcolorbox}
\vspace{-3pt}
Existing repair tools increase the number of solved tasks from 37\xspace to 42\xspace on the easy-level tasks, while increase from 5\xspace to 9\xspace on the medium-level tasks.
\vspace{-3pt}
\end{tcolorbox}
\section{RQ3: Can Codex edit mode fix program bugs?}
\label{sec:rq3}

Recently, OpenAI released a new edit mode of Codex which has the ability to change the content of an existing program. Codex edit mode takes a program and a natural language instruction as inputs, and outputs an edited program based on the instruction. As Codex-e can produce edited programs as patches, a natural question to ask would be ``Can the Codex edit mode fix an incorrect program with proper instruction?'' In this section, we explore the possibility of using Codex-e as an automated program repair tool. To reduce the influence of natural language descriptions on Codex-e, we removed the task description in each unsolved solution. We designed three strategies to construct the edit instruction for Codex-e.
\begin{itemize}[leftmargin=*]
    \item \textbf{Codex-e$^{bug}$}: We tell Codex-e that a bug exists in the given program and ask Codex-e to fix it. The instruction is simply given as ``\emph{Fix bug in the program}''.
    \item \textbf{Codex-e$^{line}$}: We follow existing automated program repair techniques that use statistical fault localization technique (Ochiai) \cite{campos2012gzoltar, ochiai} on the generated incorrect solutions to get a sequence of candidate fix line numbers. These candidate line numbers are then provided to Codex-e as fix hints. The instruction for Codex-e is formulated as ``\emph{Fix line N}''.
    \item \textbf{Codex-e$^{stm}$}: Considering the large language models like Codex are trained with plain natural language, we further investigate how Codex-e would respond if we directly use the suspicious statements instead of suspicious line numbers as instruction. To construct the edit instructions, we use the program text of the statements at the suspicious line and formulate it as ``\emph{Fix s1}''.
\end{itemize}
For example, to fix the expression mutation bug for the \textit{makeFancyString} task in Figure~\ref{listing:tbar_vs_recoder}, we give Codex-e$^{line}$ the instruction \emph{Fix line 6}, and provide Codex-e$^{stm}$ with the instruction \emph{Fix ``i -= 2;''}.

For each unsolved solution (we exclude solutions that produce syntax errors and algorithm-related error as in Section~\ref{sec:rq2}), we select the ten most suspicious statements and ask Codex-e to generate five possible edits for each statement (i.e., Codex-e tries to fix an incorrect solution within 50 attempts). Similar to the initial solution generation in the regular Codex mode, we set the temperature at 0.8 to increase the possibility of finding a correct edit.

Table~\ref{tab:data_codexe} shows the results for the three strategies, where columns Codex-e$^{bug}$, Codex-e$^{line}$ and Codex-e$^{stm}$ show the number of correct patches using corresponding edit instructions.
With \emph{Fix bug in the program} as instruction, Codex-e$^{bug}$ only learns about the existence of bugs in the given program without any information about the fault locations. Surprisingly, with the limited information, Codex-e$^{bug}$ successfully produced 15 correct patches, including two patches that need large edits at multiple hunks (M-L,  refer to supplementary material for the example).
In contrast, when giving the faulty line number as instruction, Codex-e$^{line}$ only fixes six solutions that require single-hunk fix, and one solution that require multi-hunk fixes.

Compared to Codex-e$^{bug}$ and Codex-e$^{line}$, Codex-e$^{stm}$ produces the best results by successfully fixing 16 buggy solutions.
We attribute the effectiveness of Codex-e$^{stm}$ to its use of program texts (e.g., ``i -= 2;'') that may be more helpful in guiding a language model like Codex to match for relevant statements.

\begin{figure}[!h]
\centering
\begin{lstlisting}[style=mystyle, escapechar=@]
public static int[][] construct2DArray(int[] original, int m, int n) {
  // [S-AS-3] Instruction: Fix "for (int i=0; i<result.length; i++){"
+ if (n*m != original.length)
+   return new int[0][0];
  int[][] result = new int[m][n];
  for (int i=0; i<result.length; i++){
    for (int j=0; j<result[i].length; j++){
-     if (i*result[0].length+j >= original.length)
-       return new int[0][0];
-     else
        result[i][j]=original[i*result[0].length+j];}}
  return result;
}
\end{lstlisting}
\vspace{-6pt}
\caption{Flexible fault localization example of LeetCode programming task \textit{convert-1d-array-into-2d-array} on biweekly-contest-62 fixed by Codex-e}
\label{listing:fixlocation}
\vspace{-6pt}
\end{figure}

Furthermore, we manually analyze patches produced by Codex-e, and find that Codex-e is able to generate patches at \emph{flexible locations}. Prior APR work \cite{tbar,recoder,lutellier2020coconut,jiang2021cure} have shown a significant performance gap with/without perfect fault localization results. While existing APR tools strictly try to produce patches at a given faulty line number, ignoring the possibility of fixing a bug in the relevant context, Codex-e does not have such limitations. In the 16 correctly fixed solutions by Codex-e$^{stm}$, 6 (37.5\%) of them are fixed by editing beyond the statement provided in the given instruction. Figure~\ref{listing:fixlocation} shows one such example. The instruction provided to Codex-e$^{stm}$ is \emph{Fix "for(int i =0; i<result.length; i++)\{"}, and Codex-e$^{stm}$ fix this by moving one \textit{if}-then clause out of the loop body and changing the \textit{if}-condition.

This example indicates that Codex-e$^{stm}$ is not restricted by the given suspicious statement, and can modify the relevant surrounding code. Compared to traditional APR tools, using flexible fault localization is an important feature that enables Codex-e to produce more correct patches since fault localization techniques may fail to point to the correct fix location.

\begin{tcolorbox}
\vspace{-3pt}
The efficacy of Codex-e as an automated repair tool could be worth studying. Without specific fix location, Codex-e$^{bug}$ fixes 15 buggy programs and with suspicious faulty line number and suspicious statement being provided, Codex-e$^{line}$ and Codex-e$^{stm}$ successfully fixes 7 and 16 buggy programs, respectively.

\vspace{-3pt}
\end{tcolorbox}

\section{Implications and Discussions}

\noindent \textbf{Limitations of language models.}
Apart from the mistakes common in human-crafted programs, our study in Section~\ref{sec:rq1} reveals that syntax error and misaligned algorithm are the key limitations in automatically generated solutions by Codex. Syntax errors in automatically generated programs can be solved via (1) leveraging existing techniques on compilation error repair~\cite{ahmed2018compilation,yasunaga2020graph,gupta2017deepfix}, (2) encoding a programming language model into the Codex model so that it will generate compilable programs (the approach taken by CURE~\cite{jiang2021cure}), or (3) invoking the Codex model iteratively or combining different modes of Codex to synthesize the undefined functions (the fact that we can use Codex-e for fixing the initial incorrect solutions generated by the Codex model shows the potential of this approach). The misaligned algorithm is more a severe problem that has been similarly observed by the Codex paper~\cite{codex}. Based on our manual analysis of the generated solutions, Codex seems to rely heavily on the function name (e.g., \texttt{minimumOperations} in Figure~\ref{listing:example-dp}) for solving the programming tasks. In fact, a recent study has also observed the tendency of Codex in generating solutions based on function name~\cite{jones2022capturing}. With the long prompt (function signature and the problem description) given, the function name may be more concise and easier to search for compared to the lengthy problem description in GitHub. However, this strategy fails when a customized algorithm is required to solve a programming task. Relying on function name to search for relevant code will reduce the generation power of Codex to a simple API search engine that returns the implementation for a given API.  To solve the misalignment problem, \emph{future language models designed for code generation should focus on summarizing useful information from the problem description} instead of solving the problem of generating code based on function names.

\noindent \textbf{Pattern-based APR versus learning-based APR.}

Section~\ref{sec:rq2} shows that a learning-based APR approach (Recoder) outperforms a pattern-based APR tool (TBar). We think that this result shows the limitation of a pattern-based tool that relies on pre-defined fix operators.
TBar, the state-of-art pattern-based APR tool, which integrates more than 30 specifically designed fixing patterns collected from previous literature, only fixes 6 out of 32 single-hunk buggy solutions.
Moreover, our manual analysis reveled that although TBar has the required fix patterns for some single-hunk bugs, it still fails to fix them because the concrete implementation of a fixing pattern does not always cover all possible combinations. For example, TBar has the ability to mutate a literal, variable or expression, but it does not support generating a non-existent element which is required by the \emph{makeFancyString} example shown in Figure~\ref{listing:tbar_vs_recoder}.

Given the great variety of bug combinations that Codex can make, TBar does not perform well on many buggy solutions, especially on those that require adding statement. In contrast, Recoder, a learning-based repair tool that does not rely on pre-defined fix patterns, generates a few more correct patches than TBar.
Figure~\ref{listing:tbar_vs_recoder} shows such an example, where Recoder fixes this bug by adding a statement $\mathit{steps}++;$ at line 19, which is not supported by TBar.
In short, a pattern-based approach like TBar fails to fix bugs that require either (1) additional fix patterns, or (2) a large search space for fix ingredients (e.g., specific literal). This limitations show that a \emph{pattern-based APR tool is hard to scale}. Instead of manually adding more patterns to a new APR tool, \emph{future APR research on designing fixing operators should shift to a more scalable way} (e.g., synthesizing the correct literal for the first example in Figure~\ref{listing:tbar_vs_recoder}).

\noindent \textbf{Instructions for Codex-e.} Section~\ref{sec:rq3} shows the performance of Codex-e in fixing unsolved solutions with edit instructions constructed by three different strategies. Compared to Codex-e$^{bug}$ and Codex-e$^{stm}$, Codex-e$^{line}$ generates the least number of correct fixes.
Although the number of fixed solutions by Codex-e$^{bug}$ and Codex-e$^{stm}$ are quite close (15 versus 16 bugs), the fixed defect category varies. Codex-e$^{bug}$ fixes three more multi-hunk bugs, whereas Codex-e$^{stm}$ fixes three more single-hunk bugs. Since edit instruction like \emph{Fix bug in the program} does not indicate a specific edit target, Codex-e may search for the statements to edit across the entire program based on its learned knowledge. This encourages generation of large patches but also may lose precision when fixing single-hunk bugs that only require simple fixes. In contrast, Codex-e$^{stm}$ is provided with a code context, which steers the edit to the direction that change the most relevant code context. In another perspective, we can also regard Codex-e$^{stm}$ and Codex-e$^{line}$ as test-based APR tools that fix bugs based on fault localization given by test cases, whereas Codex-e$^{bug}$ generates edit without guidance. As a test-based APR approach like Codex-e$^{stm}$ outperforms Codex-e$^{bug}$, we think that \emph{test-based APR approach is more effective than an approach without any test guidance}. Overall, \emph{patches produced by Codex-e rely heavily on the types of provided edit instruction}. Encoding the suspicious code context into the instruction performs better at fixing small bugs while general instruction without any location guidance may find more complex and larger edits. In future, it is worthwhile to study how to construct edit instructions to guide Codex-e in generating more correct fixes.

\begin{figure}[!t]
    \centering
    \includegraphics[width=0.48\textwidth]{figures/Codex-Venn.pdf}
    \caption{The repair result of different APR tools}
    \label{fig:venn_result}
\end{figure}

\noindent \textbf{Comparison between TBar, Recoder and Codex-e$^{stm}$.} To analyze the types of defects fixed by each tool and the reasons behind the effectiveness of each APR approach, we compare the patches produced by different tools. As our study in Section~\ref{sec:rq3} shows that Codex-e$^{stm}$ gives the best results among all the evaluated strategies, we select Codex-e$^{stm}$ for comparison with other APR tools.
Figure~\ref{fig:venn_result} shows a Venn diagram to better illustrate the set of commonly and uniquely produced patches by these three tools (TBar, Recoder, and Codex-e$^{stm}$). We denote the set of patches produced by TBar as $TBar$, the set of patches produced by Recoder as $Recoder$, and the set of patches produced by Codex-e$^{stm}$ as $Codex$-e$^{stm}$.
We can observe from Figure~\ref{fig:venn_result} that the set $TBar$ is a subset of $Codex$-e$^{stm}$ $\cup$ $Recoder$.

In fact, the set $TBar$ is almost subsumed by the set $Recoder$ with the only exception being the defect \emph{S-O-1} where Recoder fails to mutate a relational operator $<$ into $!=$. This is due to the limitations of pattern-based approaches discussed earlier.

If we compare Codex-e$^{stm}$ and Recoder, both approaches correctly fix 10 solutions in common, while Codex-e$^{stm}$ has eight more unique patches and Recoder has three more unique patches. We think that Codex-e$^{stm}$ outperforms Recoder because: (1) Codex-e$^{stm}$ can produce complex patches at flexible locations (as shown in the example given in Figure~\ref{listing:fixlocation}); and (2) Codex-e$^{stm}$ is trained on a much larger dataset than Recoder (Recoder uses 82868 human patches to train the repair model), which helps Codex-e$^{stm}$ to learn more fix patterns as shown in the example in Figure~\ref{listing:codex_unique} (Codex-e$^{stm}$ generates a patch that includes a lambda expression).

\begin{figure}[!h]
\centering
\begin{lstlisting}[style=mystyle, escapechar=@]
public static int minimumSum(int num){
  // S-E-1 (using lambda expresion)
- Collections.sort(digits);
+ Collections.sort(digits, (a, b) -> b - a);
}
\end{lstlisting}
\vspace{-6pt}
\caption{An example that is uniquely fixed by Codex-e$^{stm}$}
\label{listing:codex_unique}
\vspace{-6pt}
\end{figure}

Despite being trained with less data, Recoder still produces three unique patches. Figure~\ref{listing:failtofix4} shows one of the uniquely fixed solutions by Recoder. We think that Recoder can generate this correct fix due to its syntax-guided decoder that can guide it to copy the statement at line 6 and insert this statement at line 3 of Figure~\ref{listing:failtofix4} (this invokes the copy operation of Recoder that copies the AST subtree rooted at the \texttt{set.remove(i)} statement). In another example (S-E-10) uniquely fixed by Recoder, it correctly replaces a branch condition of the form \texttt{if (a \&\& b)} with \texttt{if (a)} (which is also an AST edit operation). These examples show that \emph{encoding AST information into the deep learning model may help in generating correct patches}. In future, researchers can consider incorporating AST information into large language model like Codex-e and AlphaCode to guide its patch generation.

\begin{figure}[!h]
\centering
\begin{lstlisting}[style = mystyle, escapechar=@]
public static List<List<Integer>> findWinners(int[][] matches) {  ...
  for (int i : map.keySet()) {
+   set.remove(i); // S-AS-9
    if (map.get(i) == 1) {
      ans1.add(i);
      set.remove(i);
      }  }  ...
}
\end{lstlisting}
\vspace{-6pt}
\caption{An example that is uniquely fixed by Recoder}
\label{listing:failtofix4}
\vspace{-6pt}
\end{figure}

We also find two buggy solutions (S-O-3, S-O-4) where all three APR approaches successfully generate correct fixes, but they are not efficient enough to be accepted by LeetCode{} (``Time Limit Exceeded'' error reported by LeetCode{} during submission), so we did not include them in Table~\ref{tab:correct_vs_plausible},~\ref{tab:data_apr} and ~\ref{tab:data_codexe}.
\section{Study Setting}
In this section, we present the setting of our study, including the overall workflow, the Codex model, parameters, dataset, APR tools, etc. All experiments in this paper are conducted on a Ubuntu-16.04 server, with 64GB RAM and Intel(R) Xeon(R) CPU E5-2660 v4 @ 2.00GHz, and NVIDIA Titan V GPU.

\subsection{Methodology}
Figure~\ref{fig:workflow} shows the overall workflow of our study. To answer research questions in Section~\ref{sec:intro}, we build a dataset containing \numtasks programming tasks from the recent LeetCode~ contest~\cite{leetcode}.
We first use Codex to generate initial solutions for each task and use the public tests to validate the correctness of generated solutions.
For the unsolved programming tasks, existing repair tools are then applied to the solutions produced by Codex.
The patched solutions are then validated by the given public tests and the LeetCode~online judge platform to measure how many solutions can be fixed by APR tools.

\subsection{Codex Model}
Codex~\cite{codex} is the model that powers GitHub Copilot~\footnote{https://copilot.github.com}, which completes a program by given natural language prompt.
Codex supports many programming languages (e.g., Python, C/C++, Java, Go).
The Codex models are descendants of GPT-3 models~\cite{gpt3}.
Their training data contains both natural language and billions of lines of public code from GitHub.
In our study, we use the pre-trained Codex code-davinci-002 model and Codex-e code-davinci-edit-001 model~\cite{codexmodel}, which were both trained on data up to Jun 2021.

\subsection{LMDefects{}}

Several datasets consist of programming tasks exist~\cite{alphacode,tan2017codeflaws,apps,Caballero_Description2Code_Dataset_2016,codenet,codex}. They are either based on programming contests collected from programming competition platforms such as Codeforces or hand-written tasks. We do not use existing datasets because (1) Codex has been trained on GitHub where solutions for many previous programming tasks exist (e.g., APPS, CodeContest) but we would like to check how Codex performs in a real-world programming contest setting, (2) not every programming task includes public tests which APR techniques rely on (e.g., HumanEval).

We build a new dataset LMDefects{} by mining programming tasks from LeetCode{} contest~\cite{leetcode}.
LeetCode{} is an online judge platform that people can use to practice their programming skills by solving programming tasks.
It has over 2,300 different problems, ranging from easy level to hard level.
It also has a discussion page~\footnote{https://leetcode.com/discuss/} with active community where we can find the correct solutions for a programming task (important for our manual analysis of incorrect solutions).
For each programming task, there are usually 1--3 public tests that provide examples with pairs of (input, output) to explain the requirement.
When a solution is submitted to LeetCode{}, it runs a large number of private tests to validate the correctness of submissions.
LeetCode{} has weekly and biweekly contests, where it releases new programming tasks.
Each contest includes one easy-level problem, two medium-level problems, and one hard-level problem. In our study, we only consider easy-level and the first medium-level problems, since Codex still struggles with hard problems~\cite{codex} (we also filter seven tasks that require customized data structure because those are unlikely to be solved by Codex, e.g., merge-nodes-in-between-zeros in weekly-contest-281 requires ListNode). To avoid the situation where our dataset has already been included in the training set of Codex, we only consider contests that are released after Jun 2021 (the end date where the Codex training data is extracted from).

Overall, we crawl through all contests in LeetCode from 4 July 2021 until 6 Apr 2022. This leads to a total of 40 weekly contests and 20 biweekly contests. In total, LMDefects{} contains \numtasks programming tasks, with 60 easy-level and 53 medium-level tasks, respectively.

\subsection{Prompt and Parameters}
Codex model takes as inputs a \emph{prompt}, which is the combination of natural language text and code snippet, where the natural language text represent the programming task description and code snippet is the starting point for language model to complete the code. We evaluate Codex in zero-shot prompt settings (does not include example input/output in the prompt). Note that public test cases can also be embedded in the prompt (namely \emph{n-shots prompt}), meaning that public tests are also available for Codex. Lines 1--8 in Figure~\ref{listing:example-dp} presents an example of prompt that we use for a LeetCode{} programming task~\footnote{https://leetcode.com/problems/minimum-operations-to-convert-number/}.
Given such a prompt for each programming task, we simulate the actual contest scenario by running Codex to generate 50 candidate solutions and choose the top five solutions with the highest probability of being correct (use the \emph{best\_of} parameter).
The selected solutions are then submitted to the LeetCode{} online judge platform to validate whether they are correct.
As reported in Codex~\cite{codex}, we use temperature as 0.8, which has the best performance when generating 50 candidate solutions, and we use the same setting in \cite{codex} to prepare the stop tokens as [``public'', ``class'', ``//'', ``System.out.print'']. We also follow OpenAI playground~\cite{codexmodel} default setting to set the ``max token length'' for each completion request to 256.

\begin{figure}[!t]
\centering
\begin{lstlisting}[style = mystyle, escapechar=@]
 public static int minimumOperations(int[] nums) {
/** You are given a 0-indexed array nums consisting of n positive integers.
 * The array nums is called alternating if:
 * nums[i - 2] == nums[i], where 2 <= i <= n - 1.
 * nums[i - 1] != nums[i], where 1 <= i <= n - 1.
 * In one operation, you can choose an index i and change nums[i] into any positive integer.
 * Return the minimum number of operations required to make the array alternating.
 */
  int n = nums.length;
  int[] dp = new int[n];
  dp[0] = 0; dp[1] = 1;
  for (int i = 2; i < n; i++){
    dp[i] = dp[i - 1] + 1;
    if(nums[i] == nums[i - 2])
      dp[i] = Math.min(dp[i - 2] + 1, dp[i]);
  }
  return dp[n - 1];
}
\end{lstlisting}
\vspace{-6pt}
\caption{An example of misaligned algorithm, taken from LeetCode programming task \textit{minimumOperations} on weekly-contest-280}
\label{listing:example-dp}
\vspace{-6pt}
\end{figure}

\subsection{APR Tools}

To evaluate whether repair tools can fix the incorrect solutions produced by Codex, we run our experiments on two Java APR tools. Although researchers have developed APR tools for fixing bugs in many programming languages (e.g., C, Java, Python), we use Codex to produce Java programs because Java APR tools have been widely studied, and many of them are open-source. Among all the open-source Java APR tools, we select TBar and Recoder because (1) they are the most recent representative of different approaches (i.e., TBar represents a search-based and a pattern-based APR tool, whereas Recoder is a learning-based approach), and (2) these tools have reported the best results by generating the highest number of correct patches on the Defect4J~\cite{defects4j} benchmark (a benchmark in which almost all Java APR tools has been evaluated on).

Below are the two Java repair tools used in our study:
\begin{itemize}[leftmargin=*]
    \item\textbf{TBar}: TBar~\cite{tbar} is a search-based automated program repair tool, which first constructs a search space containing a set of candidate patches, and searches for the correct patch(es) over the space space. TBar has collected a set of widely-used fix patterns from existing APR tools, including SimFix~\cite{simfix}, HDRepair~\cite{le2016history}, PAR~\cite{par2013} and so on (refer to their paper~\cite{tbar} for the complete list of work). In total, TBar supports 35 fix patterns.
    \item\textbf{Recoder}: Recoder~\cite{recoder} is a learning-based repair tool that learns from existing patches of open-source projects. Different from search-based APR tools that require predefined code change operators to generate candidate patches, learning-based tool learns code change operators from existing patches, hence has more flexibility in generating candidate patches.
\end{itemize}
Since both TBar and Recoder rely on test cases as correctness criteria, the public test cases from LeetCode are used to guide the repair process, while the private test cases are applied to validate the patched solutions. In this work, we run TBar and Recoder in default settings, stopping the repair process if one patch that passes all public tests is found. We set the timeout to 15 minutes, following the time limit used in prior work of fixing students' programming assignment using APR tools~\cite{yi2017feasibility}.

Moreover, since Codex-e can change the content of existing code by generating program edits. We investigate whether Codex-e can be further used as an APR tool and compare its performance with TBar and Recoder. We discuss how to automatically generate instructions for Codex-e at in Section~\ref{sec:rq3}.

\section{threats to validity}

\noindent\textbf{External.} During the defect categorization, one author of the paper first manually construct the ``ground truth patch'', and then discuss with the other authors to resolve any unclear categorization (e.g., when multiple fixes exist for a bug). For cases where only one fix exists for the bug, the defect classification is straightforward (we classify the defect based on the fix). As the performance of the Codex model and repair tools may varies in different settings, our experiments may not generalize beyond the studied configurations and other programming languages beyond Java. We mitigate this threat by reusing configurations given in prior work, and evaluating on several APR tools that use different algorithms (e.g., search-based and learning-based). Although other large language models (e.g., AlphaCode~\cite{alphacode}) exist, our study only evaluates on the Codex language model and the Codex edit mode. The reported findings may not generalize beyond the studied model. As the underlying algorithm used in Codex-e has not been documented, we only use it as a black-box APR tool that produces patches by editing existing programs. To ensure that the training data does not overlap with the evaluated tasks, we have confirmed with the developer of Codex-e that both Codex and Codex-e use the same dataset for training. Nevertheless, our experiments show that the Codex-e model is able to generate fixes for many unsolved solutions.

\noindent\textbf{Internal.} Our automated scripts may have bugs that can affect our reported results.
To mitigate this threat, we will make our scripts available upon acceptance.


\section{Introduction}
\label{s:intro}

In this note, we aim to engineer feature representations to aid in the
estimation of heterogeneous treatment effects. Specifically, we consider the
following graphical model
\begin{equation}
	\label{eq:model}
	\begin{tikzcd}[column sep=small]
		& X \arrow[dl] \arrow[dr] & \\
		W \arrow{rr} & & Y
	\end{tikzcd}
\end{equation}
where $X\in\ensuremath{\mathbb{R}}^d$ denotes a vector of features, $W\in\{0,1\}$ represents a
boolean treatment, and $Y\in\ensuremath{\mathbb{R}}$ denotes the outcome. Suppose $(X_i, W_i, Y_i)$
for $i=1,\dotsc,n$ are i.i.d. samples from a distribution $P$ respecting the
graph~\eqref{eq:model}. Within the potential outcomes
framework~\cite{Ney23,Rub74}, we let $Y_i(0)$ denote the potential outcome if
$W_i$ were set to $0$ and $Y_i(1)$ denote the potential outcome if $W_i$ were
set to $1$. We wish to estimate the conditional average treatment effect
({\textsc{cate}}{}) defined by
\begin{equation}
	\label{eq:cate}
	\tau(x) = \ensuremath{\mathbb{E}}[Y(1)- Y(0)|X=x].
\end{equation}
We impose standard assumptions that the treatment assignment is unconfounded,
meaning that
\[
	\{Y_i(0), Y_i(1)\} \perp W_i \mid X_i,
\]
for all $i$, and random in the sense that
\[
	\epsilon < P(W_i=1 | X_i=x_i) < 1-\epsilon
\]
for all $i$, some $\epsilon>0$, and all $x_i\in\ensuremath{\mathbb{R}}^d$ in the support of $X_i$.
These assumptions are jointly known as \emph{strong ignorability}~\cite{Ros83}
and prove sufficient for the {\textsc{cate}}{} to be identifiable. Under these
assumptions, there exist well-established methods to estimate the {\textsc{cate}}{} from
observed samples (see section~\ref{s:metalearners} for a discussion) that then
allow us to predict the expected individualized impact of an intervention for
novel examples using only their features~\cite{Pea09}. Viewing these approaches
as black box estimators, we aim to learn a mapping $\Phi: \ensuremath{\mathbb{R}}^d \to \ensuremath{\mathbb{R}}^m$ such
that the estimate of the {\textsc{cate}}{} learned from the transformed training data
$(\Phi(X_i), W_i, Y_i)$ is more accurate than an estimate learned on the
original samples $(X_i, W_i, Y_i)$. In particular, we desire a function $\Phi$
yielding a corresponding representation $\Phi(X)$ such that
\begin{enumerate}
	\item $\Phi(X)$ is as useful as $X$ for estimating $Y$, and
	\item among such representations, $\Phi(X)$ is least useful for estimating
	      $W$.
\end{enumerate}
In this way, we hope to produce a new set of features $\Phi(X)$ that retain all
information relevant for predicting the outcome, but are less related to
treatment assignment. \emph{We propose learning $\Phi$ as an intermediate layer
	in a neural network estimating a functional relationship of $Y$ given $X$.
	We apply a genetic algorithm~\cite{Hol75} to a population of such mappings
	to evolve and select the one for which the associated representation $\Phi(X)$
	is least useful for approximating $W$.}

Feature representations are commonly used in machine learning to aid in the
training of supervised models~\cite{Ben13} and have been previously
demonstrated to aid in causal modeling. Johansson, et al.~\cite{Joh16,Sha17b}
viewed counterfactual inference on observational data as a covariate shift
problem and learned neural network-based representations designed to produce
similar empirical distributions among the treatment and control populations,
namely $\{\Phi(X_i)\}_{W_i=1}$ and $\{\Phi(X_i)\}_{W_i=0}$. Li \&
Fu~\cite{Li17a} and Yao, et al.~\cite{Yao18} developed representations in a
related vein designed to preserve local similarity. We generally agree with
Zhang et al.'s~\cite{Zha20} recent argument that domain invariance
often removes too much information from the features for causal
inference.\footnote[2]{Zhao et al.~\cite{Zha19a} make this argument in a more
	general setting.} \emph{In contrast to most previous approaches, we develop a
	feature representation that attempts to preserve information useful for
	predicting the treatment effect if it is also useful for predicting the
	outcome.}

\subsubsection{Outline.}
We proceed as follows. In the next section, we describe methods for learning
the {\textsc{cate}}{} from observational data and introduce genetic algorithms. In
section~\ref{s:methodology}, we describe our methodology in full. We then
validate our method on artificial data in section~\ref{s:ablation_synthetic}
and on a publicly available experimental dataset in
section~\ref{s:in_house_data}, before concluding in
section~\ref{s:conclusions}.

\section{Related work}
\label{s:related}

In the first part of this section, we discuss standard methods for learning the
{\textsc{cate}}{} function from data. We will subsequently use these to test our proposed
feature engineering methods in section~\ref{s:ablation_synthetic}. In the
second part, we briefly outline evolutionary algorithms for training neural
networks, commonly called neuroevolutionary methods.

\subsection{Meta-learners}
\label{s:metalearners}
We adopt the standard assumptions of unconfoundedness and the random assignment
of treatment effects that together constitute strong ignorability. Given i.i.d.
samples from a distribution $P$ respecting~\eqref{eq:model} and these
assumptions, there exist numerous meta-learning approaches that leverage an
arbitrary regression framework (e.g., random forests, neural networks,
linear regression models) to estimate the {\textsc{cate}}{} that we now describe.

\subsubsection*{S-learner.}
The S-learner (single-learner) uses a standard supervised learner
(regression model) to estimate $\mu(x,w) = \ensuremath{\mathbb{E}}[Y | X=x, W=w]$ from observation
data and then predicts
\[
	\hat\tau_S(x) = \hat\mu(x,1) - \hat\mu(x,0)
\]
where we use the standard hat notation to denote estimated versions of the
underlying functions.

\subsubsection*{T-learner.}
The T-learner (two-learner) estimates $\mu_1(x) = \ensuremath{\mathbb{E}}[Y(1)|X=x]$ from observed
treatment data $\{(X_i,Y_i)\}_{W_i=1}$ and $\mu_0(x) = \ensuremath{\mathbb{E}}[Y(0)|X=x]$ from
observed control data $\{(X_i,Y_i)\}_{W_i=0}$, and then predicts
\[
	\hat\tau_T(x) = \hat\mu_1(x) - \hat\mu_0(x).
\]

\subsubsection*{X-learner.}
The X-learner~\cite{Kun19} estimates $\mu_1$ and $\mu_0$ as in the T-learner,
and then predicts the contrapositive outcome for each training point. Next,
the algorithm estimates $\tau_1(x)=\ensuremath{\mathbb{E}}[\tilde D_i^1 \mid X=x]$ on
$\{(X_i,\tilde D_i^1) \}_{W_i=1}$ where $\tilde D_i^1 = Y_i - \hat\mu_0(X_i)$
and $\tau_0(x)=\ensuremath{\mathbb{E}}[\tilde D_i^0 \mid X=x]$ on $\{(X_i,\tilde D_i^0)
	\}_{W_i=0}$ where $\tilde D_i^0 = \hat\mu_1(X_i)-Y_i$. The X-learner then
predicts
\[
	\hat\tau_X(x) = g(x)\hat\tau_0(x) + (1-g(x))\hat\tau_1(x)
\]
where $g: \ensuremath{\mathbb{R}}^d\to [0,1]$ is a weight function. The creators of the X-learner
remark that the treatment propensity function~\eqref{eq:treatment_propensity}
often works well for $g$, as do the constant functions $1$ and $0$. In our
implementation, we use $g(x)\equiv 1/2$.

Concerning this method, we note that it is also possible to directly estimate
$\tau$ from
\[
	\{(X_i,Y_i - \hat\mu_0(X_i)) \}_{W_i=1}
	\cup \{(X_i,\hat\mu_1(X_i)-Y_i) \}_{W_i=0}
\]
or, using $\hat\mu(x,w)$ from the S-learner approach, with
\[
	\{(X_i,Y_i - \hat\mu(X_i,0)) \}_{W_i=1}
	\cup \{(X_i,\hat\mu(X_i,1)-Y_i) \}_{W_i=0}.
\]
We find that these alternate approaches work well in practice and obviate the
need to estimate or fix $g$.

\subsubsection*{R-learner.}
Within the setting of the graphical model~\eqref{eq:model}, we define the
treatment propensity (sometimes called the propensity score) as
\begin{equation}
	\label{eq:treatment_propensity}
	e(x) = P(W=1 | X=x)
\end{equation}
and the conditional mean outcome as
\begin{equation}
	\label{eq:conditional_mean_outcome}
	m(x) = \ensuremath{\mathbb{E}}[Y | X=x].
\end{equation}
The R-learner~\cite{Nie21} leverages Robinson's decomposition~\cite{Rob88} that
led to Robin's reformulation~\cite{Rob04} of the {\textsc{cate}}{} function as the
solution to the optimization problem
\begin{equation}
	\label{eq:cate_as_optimization}
	\tau(\cdot) = \argmin_{\tau} \bigg\{ \ensuremath{\mathbb{E}}_{(X,W,Y)\sim P}
	\bigg[ \abs{\big(Y - m(X)\big) - \big(W - e(X)\big)\tau(X)}^2 \bigg] \bigg\}
\end{equation}
in terms of the treatment propensity~\eqref{eq:treatment_propensity} and the
conditional mean outcome~\eqref{eq:conditional_mean_outcome}. In practice, a
regularized, empirical version of~\eqref{eq:cate_as_optimization} is minimized
via a two-step process: (1) cross-validated estimates $\hat m$ and $\hat e$ are
obtained for $m$ and $e$, respectively, and then (2) the empirical loss is
evaluated using folds of the data not used for estimating $\hat m$ and $\hat
	e$, and then minimized. The authors Nie \& Wager note that the structure of the
loss function eliminates correlations between $m$ and $e$ while allowing one to
separately specify the form of $\tau$ through the choice of optimization
method. In this paper, the only R-learner we use is the causal forest as
implemented with generalized random forests~\cite{Ath19} using the default
options, including honest splitting~\cite{Wag18}.

\subsection{Genetic and neuroevolutionary algorithms}
Holland introduced genetic algorithms~\cite{Hol75} as a nature-inspired
approach to optimization. Generally speaking, these algorithms produce
successive generations of candidate solutions. New generations are formed by
selecting the fittest members from the previous generation and performing
cross-over and/or mutation operations on them to produce new offspring
candidates. Evolutionary algorithms encompass extensions and generalizations to
this approach including memetic algorithms~\cite{Mos89} that perform local
refinements, genetic programming~\cite{For81} that acts on programs represented
as trees, and evolutionary programming~\cite{Fog62} and
strategies~\cite{Rec70,Sch75} that operate on more general representations.
When such methods are applied specifically to the design and training of neural
networks, they are commonly known as neuroevolutionary algorithms. See Stanley
et al.~\cite{Sta19} for a comprehensive survey. In the next section, we
describe a specific neuroevolutionary strategy for feature engineering.

\section{Methodology}
\label{s:methodology}

In this section, we describe how we form our feature mapping
$\Phi:\ensuremath{\mathbb{R}}^d\to\ensuremath{\mathbb{R}}^m$. To generate a single candidate solution, we train a
shallow neural network to predict $Y$ from $X$ and extract an intermediate
layer of this network. Each candidate map created in this way should yield a
representation as functionally useful for predicting $Y$ as $X$ is. We then
iteratively evolve cohorts of parameter sets for such maps to create a
representation that carries the least amount of useful information for
predicting the treatment $W$.

\subsection{Candidate solutions}
We consider neural networks $f_\Theta: \ensuremath{\mathbb{R}}^d \to \ensuremath{\mathbb{R}}$ of the form
\begin{equation}
	\label{eq:f_theta}
	f_\Theta(x) = M_2 \cdot a( M_1 \cdot x + b_1 ) + b_2
\end{equation}
where $M_1\in \ensuremath{\mathbb{R}}^{m \times d}$, $M_2 \in \ensuremath{\mathbb{R}}^{1\times m}$ are real-valued
matrices (often called weights), $b_1\in \ensuremath{\mathbb{R}}^m$ and $b_2 \in \ensuremath{\mathbb{R}}^1$ are vectors
(often called biases), and $a$ is a nonlinear activation function applied
component-wise. We let $\Theta = ( M_1, M_2, b_1, b_2)$ denote the parameters
for $f_\Theta$. Though $f_\Theta$ is decidedly not a deep neural network, we
note that, as a neural network with a single hidden layer, it remains a
universal function approximator in the sense of Hornik et al.~\cite{Hor89}.
Optimizing the network~\eqref{eq:f_theta} in order to best predict $Y$ from $X$
seeks the solution
\begin{equation}
	\label{eq:optimize_f}
	\Theta_* = \argmin_\Theta \ensuremath{\mathbb{E}} \abs{ Y - f_\Theta(X) }^2.
\end{equation}
Given parameters $\Theta$ for the network~\eqref{eq:f_theta}, we let
$\Phi_\Theta: \ensuremath{\mathbb{R}}^d\to\ensuremath{\mathbb{R}}^m$ given by
\begin{equation}
	\label{eq:phi_theta}
	\Phi_\Theta(x) = a( M_1 \cdot x + b_1 )
\end{equation}
denote the output of the hidden layer. We restrict to candidate feature
mappings of this form. As these mappings are completely characterized by their
associated parameters, we define a fitness function and evolutionary algorithm
directly in terms of parameter sets $\Theta$ in the following subsections.

\subsection{Fitness function}
For parameters $\Theta$ near the optimum~\eqref{eq:optimize_f}, we note that
$\Phi_\Theta(X)$ should be approximately as useful as $X$ for learning a
functional relationship with $Y$. However, for some values of $\Theta$, the
mapped features $\Phi_\Theta(X)$ may carry information useful for predicting
$W$, and for this reason we consider a network $g_{\Psi,\Theta}: \ensuremath{\mathbb{R}}^d\to[0,1]$
given by
\begin{equation}
	\label{eq:g_psi}
	g_{\Psi,\Theta}(x)
	= \sigma( M_4 \cdot a (M_3 \cdot \Phi_\Theta(x) + b_3) + b_4)
\end{equation}
where $M_3 \in \ensuremath{\mathbb{R}}^{k \times m}$ and $M_4 \in \ensuremath{\mathbb{R}}^{1 \times k}$ are weights,
$b_3\in \ensuremath{\mathbb{R}}^k$ and $b_4 \in \ensuremath{\mathbb{R}}$ are biases, $a$ is a nonlinear activation
function applied component-wise,\footnote{though not necessarily the same as
	the one used in~\eqref{eq:f_theta} and \eqref{eq:phi_theta}} and $\sigma(x) =
	(1+\exp(-x))^{-1}$ denotes the sigmoidal activation function. In this case,
$\Psi = ( M_3, M_4, b_3, b_4)$ denotes the collection of tunable parameters. We
define the fitness of a parameter set $\Theta$ to be
\begin{equation}
	\label{eq:mu_theta}
	\mu(\Theta) = \min_{\Psi} \ensuremath{\mathbb{E}} \abs{ W - g_{\Psi,\Theta}(X)}^2.
\end{equation}
In this way, we express a preference for representations $\Phi_\Theta(X)$ that
are less useful for predicting $W$. For a schematic of these architectures,
please see \figurename~\ref{f:schematic}.

\begin{figure}[htb]
	\centering
	\begin{minipage}[b]{0.4\textwidth}
		\centering
		\includegraphics[width=\textwidth]{nn_schematic.pdf}
	\end{minipage}
	\hfill
	\begin{minipage}[b]{0.55\textwidth}
		\caption{In this schematic, the arrows connecting the $X_i$ to the
			$\Phi(X)_j$ represent the map $\Phi_\Theta$ as
			in~\eqref{eq:phi_theta}; these, in addition to the ones joining the
			$\Phi(X)_j$ to $Y$, represent $f_\Theta$ given explicitly
			in~\eqref{eq:f_theta}; and the original arrows along with the
			dashed arrows represent $g_{\Psi,\Theta}$ as in~\eqref{eq:g_psi}.}
		\label{f:schematic}
	\end{minipage}
\end{figure}

\subsection{Evolutionary algorithm}
We now describe a method to generate and evolve a cohort of candidate parameter
sets $\Theta$ intended to seek a parameter set $\Theta_*$ such that
$\Phi_{\Theta*}(X)$ is nearly as useful for predicting $Y$ as $X$ is and, among
such representations, $\Phi_{\Theta*}(X)$ is least useful for predicting $W$.

Given training data $(X_i, W_i, Y_i) \sim^\text{i.i.d.} P$ for $1\leq i \leq
	n$, we first partition the data into training and validation sets. We form an
initial cohort of $c$ candidates independently as follows. For $1\leq j \leq
	c$, we randomly instantiate $\Theta_j$ using Glorot normal
initialization~\cite{Glo10} for the weights and zeros for the biases and then
apply batch-based gradient descent on the training set to seek the
solution to~\eqref{eq:optimize_f}. In particular, we use the Adam
optimizer~\cite{Kin15} that maintains parameter-specific learning
rates~\cite[cf. AdaGrad]{Duc11} and allows these rates to sometimes
increase~\cite[cf. Adadelta]{Zei12} by adapting them using the first two
moments from recent gradient updates. We use Tikhonov
regularization~\cite{Tik63} for the weights and apply a dropout
layer~\cite{Sri14} after the $a(x) = \tanh(x)$ activation function\footnote[4]{
	We tested rectified~\cite{Nai10} and exponential~\cite{Cle16} linear unit
	activation functions for $a$ in $\Phi_\Theta$ but noticed only minor
	differences in subsequent performance of the causal forest.}%
to prevent overfitting.

For each constituent $\Theta_j$ in the cohort, we then initialize and train a
network $g_{\Psi,\Theta_j}$ as in~\eqref{eq:g_psi} to seek $\Psi_j =
	\argmin_{\Psi} \ensuremath{\mathbb{E}} \abs{ W - g_{\Psi,\Theta_j}(X)}^2$ on the training set and
then evaluate $\ensuremath{\mathbb{E}} \abs{ W - g_{\Psi_j,\Theta_j}(X)}^2$ empirically on the
validation set to estimate $\mu(\Theta_j)$. We then use the $\ell$ fittest
members of the current cohort to form a new cohort as follows. For each of the
$\binom{\ell}{2}$ pairings, we apply Montana and Davis's node-based
crossover~\cite{Mon89} method to the parameters $M_1$ and $b_1$ that we use to
form $\Phi$. This amounts to forming a new $\Phi$ by randomly selecting one of
the two parents and using that parent's mapping for each coordinate. Thus, the
new $M_1$ and $b_1$ are selected in a row-wise manner from the corresponding
rows of the parents, and then the new $M_2$ and $b_2$ are randomly initialized
and a few steps of optimization are performed to form the offspring candidate.
The next generation then consists of the best performing candidate from the
previous generation, $\binom{\ell}{2}$~candidates formed by crossing the best
$\ell$ candidates of the previous generation, and $c-1-\binom{\ell}{2}$
entirely new candidates generated from scratch.

\begin{algorithm}[!tb]
	\SetAlgoLined
	\KwData{training $\mathcal{T} = \{(X_i, W_i, Y_i)\}_{i \in T}$ and
	validation $\mathcal{V} = \{(X_i, W_i, Y_i)\}_{i \in V}$ datasets drawn
	i.i.d. from $P$ respecting the graphical model~\eqref{eq:model}\; positive
	integer parameters: $c$ cohort size, $\ell$ number of members involved in
	forming the next generation, $g$ number of generations, $m$ dimensionality
	of the latent representation, and $k$ dimensionality of the final hidden
	layer in~\eqref{eq:g_psi}}
	\KwResult{parameterized function $\Phi_{\Theta_*}: \ensuremath{\mathbb{R}}^d \to \ensuremath{\mathbb{R}}^m$ such
	that models for the {\textsc{cate}}{} learned using the transformed training data
	$\{(\Phi(X_i), W_i, Y_i)\}_{i\in \mathcal{T}}$ perform better than those
	learned on the original dataset $\mathcal{T}$}
	\For{$j = 1,\dotsc,c$}{ Optimize a parameter set $\Theta_j$ to seek
		\eqref{eq:optimize_f} on training batches from a random initialization\;
	}
	Form the first generation $\mathcal G_1 = \{ \Theta_j \}_{j = 1,\dotsc, c}$\;
	\For(\#\emph{form the next generation}){$t = 2,\dotsc,g$}{ Initialize
		new generation $\mathcal G_t = \{ \argmax_{\Theta\in\mathcal{G}_{t-1}}
			\mu(\Theta) \}$ with the best-performing candidate from the previous
		generation\;
		\For(\#\emph{form new candidates using crossover}){unique pairs
			$\{\Theta_j, \Theta_k\}$ formed from the top $\ell$ candidates from
			$\mathcal{G}_{t-1}$}{
			Initialize $M_1\in \ensuremath{\mathbb{R}}^{m\times d}$ and $b_1\in\ensuremath{\mathbb{R}}^m$ as the $M_1$ and
			$b_1$ from $\Theta_j$\; \For{$ \kappa= 1,\dotsc,m$}{ Let $\xi\sim
					\operatorname{Bernoulli}(1/2)$\; \lIf{$\xi = 1$}{
					replace the $\kappa$th row of $M_1$ and the $\kappa$th component of
					$b_1$ with those from $\Theta_k$}
			}
			Randomly initialize $M_2$ and $b_2$ and take optimization steps towards
			the solution to~\eqref{eq:optimize_f}\; Add $\Theta = ( M_1, M_2, b_1,
				b_2)$ to $\mathcal G_t$
		}
		\While{$\abs{\mathcal G_t}<c$}{ Optimize a parameter set $\Theta$ to seek
			\eqref{eq:optimize_f} on training batches from a random initialization\;
			Add $\Theta$ to $\mathcal G_t$
		}
	}
	Let $\Theta_* = \argmax_{\Theta\in\mathcal{G}_g} \mu(\Theta)$\;
	\Return{$\Phi_{\Theta_*}$ as in \eqref{eq:phi_theta}}
	\caption{Neuroevolutionary Feature Engineering for Causal Inference}
	\label{algo:main}
\end{algorithm}

We summarize our approach using pseudo-code in
\algorithmcfname~\ref{algo:main}. All computation of the valuation function is
done by training a network of the form~\eqref{eq:g_psi} to minimize $\abs{ W -
		g_{\Psi,\Theta}(X)}^2$ on training batches and then
approximating~\eqref{eq:mu_theta} by taking the empirical mean on the
validation set.

\subsection{Remark on linearity}
Due to our choice of representation $\Phi$ in~\eqref{eq:phi_theta}, after
training the network~\eqref{eq:f_theta} to optimize~\eqref{eq:optimize_f}, we
expect the relationship between the learned features $\Phi(X)$ and the outcome
$Y$ to be approximately linear. In particular, we will have $Y\approx M_2\cdot
	\Phi(X)$ for $M_2$ as given in~\eqref{eq:f_theta}. For this reason, the
causal meta-learners trained using a linear regression base learner may
benefit more extensively from using the transformed features instead of the
original features, especially in cases where the relationship between the
original features and outcomes is not well-approximated as linear. We provide
specific examples in the next section involving meta-learners trained with
ridge regression.

\subsection{Remark on our assumptions}
In order to use the represented features $\Phi(X_i)$ in place of the original
features $X_i$ when learning the {\textsc{cate}}{}, we require that strong ignorability
holds for the transformed dataset $(\Phi(X_i), W_i, Y_i)$, $i=1,\dotsc,n$. One
sufficient, though generally not necessary, assumption that would imply strong
ignorability is for $\Phi$ to be invertible on the support of
$X$~\cite[assumption 1]{Sha17b}. Unconfoundedness would also be guaranteed if
$\Phi(X)$ satisfied the backdoor condition with respect to
$(W,Y)$~\cite[section 3.3.1]{Pea09}.

\section{Ablation study on generated data}
\label{s:ablation_synthetic}

Due to the fundamental challenge of causal inference (namely, that the
counterfactual outcome cannot be observed, even in controlled experiments), it
is common practice to compare approaches to {\textsc{cate}}{} estimation on artificially
generated datasets that allow the {\textsc{cate}}{} to be calculated directly for
evaluative purposes. In this section, we perform experiments using two such
data generation mechanisms from Nie \& Wager's paper~\cite{Nie21} that we now
describe.\footnote[3]{Nie \& Wager's paper included four setups,
	namely A--D; however setup B modeled a controlled randomized trial and setup D
	had unrelated treatment and control arms.} Both setups provide a joint
distribution satisfying the graph~\eqref{eq:model}. For the vector-valued
random variable $X\in\ensuremath{\mathbb{R}}^d$, we let $X_{ij}$ denote $j$th component ($1\leq j
	\leq d$) of the $i$th sample ($1\leq i \leq n$). The specifics for both setups
are given as follows.

\subsubsection{Setup A.}
For $\sigma>0$ and an integer $d>0$, we let
\[
	X_i\sim^\text{i.i.d.} \operatorname{Uniform}([0,1]^d)
	\quad\text{and}\quad
	W_i \mid X_i \sim \operatorname{Bernoulli}(e(X_i)),
\]
where $e(X_i) = \max\{0.1, \min\{\sin(\pi X_{i1}X_{i2}), 0.9\} \}$ and
\[
	Y_i \mid X_i,W_i
	\sim \mathcal{N}\big(b(X_i) + (W_i-0.5)\tau(X_i),\sigma^2\big),
\]
where $b(X_i) = \sin(\pi X_{i1}X_{i2}) + 2 (X_{i3}-0.5)^2 + X_{i4} + 0.5
	X_{i5}$ and $\tau(X_i) = (X_{i1}+X_{i2}) /2$. In this paper, we let $d=24$,
$n=200$, and $\sigma=1$.

\subsubsection{Setup C.}
For $\sigma>0$ and an integer $d>0$, we let
\[
	X_i\sim^\text{i.i.d.} \mathcal{N}_d(\vec 0, I_{d\times d})
	\quad\text{and}\quad
	W_i \mid X_i \sim \operatorname{Bernoulli}(e(X_i)),
\]
where in this case $e(X_i) = (1+ \exp(X_{i2} + X_{i3}) )^{-1}$ and
\[
	Y_i \mid X_i, W_i
	\sim \mathcal{N}\big( b(X_i) + (W_i-0.5) \tau(X_i), \sigma^2\big),
\]
where now $b(X_i) = 2 \log(1+ \exp(X_{i1} + X_{i2} + X_{i3}))$
and $\tau(X_i) = 1$. In our example, we let $d=12$, $n=500$, and $\sigma=1$.

\subsection{Comparison methodology}
For each data generation method, we ran 100 independent trials. Within each
trial, we simulated a dataset of size $n$ and randomly partitioned it into
training, validation, and testing subsets at a 70\%-15\%-15\% rate. We trained
causal inference methods on the training set, using the validation data to aid
the training of some base estimators for the meta-learners, and predicted on
the test dataset. We then developed a feature map using the training and
validation data as described in the previous section, applied this map to all
features, and repeated the training and testing process using the new features.

To determine the impact of the fitness selection process, we also learned a
feature transformation that did not make use of the fitness function at all. In
effect, it simply generated a single candidate mapping and used it to transform
all the features (without any cross-over or further mutation). Features
developed in this way are described as ``no fitness'' in the tables that
follow.

We compared the causal forest with default options (as found in R's grf
package), and the S-, T-, and X-learners as described in
section~\ref{s:metalearners} with two base learners. The first base learner is
LightGBM~\cite{Ke17}, a boosted random forest algorithm that introduced novel
techniques for sampling and feature bundling. The second is a cross-validated
ridge regression model (as found in scikit-learn) that performs multiple linear
regression with an $L^2$ normalization on the weights.

\subsection{Results}

We report results for setup A in \tablename~\ref{t:setup_A_results} and results
for setup C in \tablename~\ref{t:setup_C_results}. For both setups, we
consider a paired t-test for equal means against a two-sided alternative (as
implemented in Python's scipy package). For setup A, we find that the
improvement in MSE from using the transformed features in place of the original
features corresponds to a statistically significant difference for the
following learners: the causal forest ($p<0.001$), the S-learner with ridge
regression ($p<0.001$), the T-learner with both LightGBM ($p<0.001$) and ridge
regression ($p<0.001$), and the X-learner with both LightGBM ($p<0.001$) and
ridge regression ($p<0.001$). For setup C, we again find significant
differences for the causal forest ($p=0.023$), S-learner with LightGBM
($p=0.003$), T-learner with ridge regression ($p<0.001$) and X-learner with
ridge regression ($p<0.001$).

In summary, we find that our feature transformation method improves the
performance of multiple standard estimators for the {\textsc{cate}}{} under two data
generation models.

\begin{table}[htb]
	\centering
	\begin{tabular}{llccccccc>{\bfseries}c}
		\toprule
		learner       & features    & tr.\,1 & tr.\,2 & tr.\,3 & \dots & tr.\,98 & tr.\,99 & tr.\,100 & avg.  \\ \midrule
		Causal forest & initial     & 0.488  & 0.161  & 0.042  & \dots & 0.207   & 0.212   & 0.592    & 0.175 \\
		              & no fitness  & 0.305  & 0.050  & 0.035  & \dots & 0.048   & 0.068   & 0.353    & 0.114 \\
		              & transformed & 0.163  & 0.054  & 0.044  & \dots & 0.046   & 0.103   & 0.386    & 0.120 \\ \midrule
		S-L. w/ LGBM  & initial     & 0.121  & 0.055  & 0.191  & \dots & 0.048   & 0.059   & 0.323    & 0.149 \\
		              & no fitness  & 0.171  & 0.111  & 0.089  & \dots & 0.149   & 0.098   & 0.395    & 0.140 \\
		              & transformed & 0.157  & 0.077  & 0.096  & \dots & 0.168   & 0.176   & 0.241    & 0.135 \\ \midrule
		S-L. w/ Ridge & initial     & 0.257  & 0.038  & 0.032  & \dots & 0.048   & 0.067   & 0.290    & 0.093 \\
		              & no fitness  & 0.202  & 0.038  & 0.037  & \dots & 0.039   & 0.046   & 0.249    & 0.079 \\
		              & transformed & 0.110  & 0.039  & 0.051  & \dots & 0.045   & 0.041   & 0.274    & 0.081 \\ \midrule
		T-L. w/ LGBM  & initial     & 0.574  & 0.351  & 0.691  & \dots & 0.448   & 0.198   & 1.189    & 0.666 \\
		              & no fitness  & 0.647  & 0.516  & 0.596  & \dots & 0.393   & 0.303   & 0.669    & 0.536 \\
		              & transformed & 0.469  & 0.295  & 0.599  & \dots & 0.334   & 0.504   & 1.056    & 0.512 \\ \midrule
		T-L. w/ Ridge & initial     & 0.763  & 0.254  & 0.888  & \dots & 0.916   & 0.260   & 0.553    & 0.745 \\
		              & no fitness  & 0.257  & 0.177  & 0.103  & \dots & 0.237   & 0.114   & 0.574    & 0.333 \\
		              & transformed & 0.147  & 0.042  & 0.206  & \dots & 0.087   & 0.118   & 0.471    & 0.325 \\ \midrule
		X-L. w/ LGBM  & initial     & 0.339  & 0.207  & 0.309  & \dots & 0.394   & 0.134   & 0.811    & 0.411 \\
		              & no fitness  & 0.497  & 0.225  & 0.230  & \dots & 0.161   & 0.128   & 0.493    & 0.335 \\
		              & transformed & 0.361  & 0.135  & 0.317  & \dots & 0.124   & 0.368   & 0.757    & 0.317 \\ \midrule
		X-L. w/ Ridge & initial     & 0.643  & 0.221  & 0.702  & \dots & 0.798   & 0.240   & 0.470    & 0.630 \\
		              & no fitness  & 0.247  & 0.145  & 0.097  & \dots & 0.197   & 0.110   & 0.540    & 0.289 \\
		              & transformed & 0.149  & 0.033  & 0.176  & \dots & 0.066   & 0.118   & 0.460    & 0.288 \\
		\bottomrule
	\end{tabular}
	\vspace{4pt}
	\caption{Mean Squared Error (MSE) over 100 independent trials run using
		setup~A. Algorithm~\ref{algo:main} was run with parameters: cohort size
		$c=4$, progenitors $\ell=2$, number of cohorts $g=5$, representation
		dimensionality $m=20$, and fitness function parameter $k=10$.}
	\label{t:setup_A_results}
\end{table}

\begin{table}[htb]
	\centering
	\begin{tabular}{llccccccc>{\bfseries}c}
		\toprule
		learner       & features    & tr.\,1 & tr.\,2 & tr.\,3 & \dots & tr.\,98 & tr.\,99 & tr.\,100 & avg.  \\ \midrule
		Causal forest & initial     & 0.002  & 0.062  & 0.016  & \dots & 0.023   & 0.008   & 0.071    & 0.035 \\
		              & no fitness  & 0.004  & 0.027  & 0.028  & \dots & 0.046   & 0.006   & 0.047    & 0.029 \\
		              & transformed & 0.007  & 0.024  & 0.028  & \dots & 0.045   & 0.003   & 0.034    & 0.029 \\ \midrule
		S-L. w/ LGBM  & initial     & 0.247  & 0.244  & 0.279  & \dots & 0.222   & 0.190   & 0.202    & 0.226 \\
		              & no fitness  & 0.173  & 0.207  & 0.207  & \dots & 0.214   & 0.164   & 0.198    & 0.211 \\
		              & transformed & 0.188  & 0.183  & 0.146  & \dots & 0.331   & 0.149   & 0.196    & 0.204 \\ \midrule
		S-L. w/ Ridge & initial     & 0.000  & 0.024  & 0.017  & \dots & 0.004   & 0.012   & 0.005    & 0.015 \\
		              & no fitness  & 0.000  & 0.016  & 0.017  & \dots & 0.005   & 0.012   & 0.005    & 0.015 \\
		              & transformed & 0.000  & 0.022  & 0.018  & \dots & 0.005   & 0.025   & 0.006    & 0.015 \\ \midrule
		T-L. w/ LGBM  & initial     & 0.771  & 0.482  & 0.833  & \dots & 0.490   & 0.538   & 0.424    & 0.567 \\
		              & no fitness  & 0.587  & 0.543  & 0.593  & \dots & 0.491   & 0.766   & 0.537    & 0.551 \\
		              & transformed & 0.353  & 0.541  & 0.378  & \dots & 0.551   & 0.575   & 0.531    & 0.544 \\ \midrule
		T-L. w/ Ridge & initial     & 0.143  & 0.260  & 0.189  & \dots & 0.173   & 0.148   & 0.124    & 0.178 \\
		              & no fitness  & 0.049  & 0.164  & 0.096  & \dots & 0.145   & 0.106   & 0.093    & 0.125 \\
		              & transformed & 0.084  & 0.136  & 0.089  & \dots & 0.116   & 0.095   & 0.109    & 0.128 \\ \midrule
		X-L. w/ LGBM  & initial     & 0.286  & 0.216  & 0.474  & \dots & 0.282   & 0.247   & 0.275    & 0.313 \\
		              & no fitness  & 0.348  & 0.236  & 0.361  & \dots & 0.318   & 0.377   & 0.254    & 0.313 \\
		              & transformed & 0.342  & 0.193  & 0.207  & \dots & 0.208   & 0.278   & 0.231    & 0.298 \\ \midrule
		X-L. w/ Ridge & initial     & 0.129  & 0.241  & 0.170  & \dots & 0.156   & 0.136   & 0.114    & 0.166 \\
		              & no fitness  & 0.039  & 0.142  & 0.081  & \dots & 0.119   & 0.099   & 0.079    & 0.109 \\
		              & transformed & 0.072  & 0.124  & 0.078  & \dots & 0.104   & 0.082   & 0.097    & 0.114 \\
		\bottomrule
	\end{tabular}
	\vspace{4pt}
	\caption{Mean Squared Error (MSE) over 100 independent trials run using
		setup~C. Algorithm~\ref{algo:main} was run with
		$c=4$, $\ell=2$, $g=2$, $m=20$, and $k=10$.}
	\label{t:setup_C_results}
\end{table}

\section{Application to econometric data}
\label{s:in_house_data}

In this section, we apply our feature engineering method to the LaLonde
dataset~\cite{LaL86,Deh99} chronicling the results of an experimental study on
temporary employment opportunities. The dataset contains information from 445
participants who were randomly assigned to either an experimental group that
received a temporary job and career counseling or to a control group that
received no assistance. Features include age and education (in years), earnings
in 1974 (in \$, prior to treatment), and indicators for African-American
heritage, Hispanic-American heritage, marital status, and possession of a high
school diploma. We consider the outcome of earnings in 1978 (in \$, after
treatment).

We cannot determine true average treatment effects based on individual-level
characteristics (i.e. the true {\textsc{cate}}{} values) for real life experimental data
as we can with the synthetic examples of the previous section. Instead, we
evaluate performance by comparing the average realized treatment effect and
average predicted treatment effect within bins formed by sorting study
participants according to predicted treatment effect as demonstrated in
\figurename~\ref{fig:lalonde}. Applying the causal forest predictor to the
original features results in a root mean square difference between the average
predicted and realized treatment effects of 4729.51. If the transformed
features are used instead, this discrepancy improves to 3114.82.

From a practical perspective, one may learn the {\textsc{cate}}{} in order to select a
subset of people for whom a given intervention has an expected net benefit (and
then deliver that intervention only to persons predicted to benefit from it).
When we focus on the 20\% of people predicted to benefit most from this
treatment, we find that the estimated realized benefit for those chosen using
the transformed features (\$4732.89) is much greater than the benefit for those
chosen using the original feature set (\$816.92). This can be seen visually in
\figurename~\ref{fig:lalonde} by comparing the estimated realized average
treatment effect for bin \#5 (the rightmost bin) in both plots.

\begin{figure}[tb]
	\centering
	\begin{subfigure}[b]{0.475\textwidth}
		\centering
		\includegraphics[width=\textwidth]{lalonde_initial.pdf}
	\end{subfigure}
	\hspace{10pt}
	\begin{subfigure}[b]{0.475\textwidth}
		\centering
		\includegraphics[width=\textwidth]{lalonde_transformed.pdf}
	\end{subfigure}
	\caption{We plot estimated realized and predicted average treatment effects
		versus the quintiles of predicted treatment effect (5 bins) for a
		causal forest using (left) the initial features and (right) the features
		transformed using our method. In order to estimate realized treatment
		effect within each bin, we take the difference between the average
		outcome for persons randomly assigned to the experimental group and the
		average outcome for persons randomly assigned to the control group.}
	\label{fig:lalonde}
\end{figure}

\section{Conclusions}
\label{s:conclusions}

Causal inference, especially on real life datasets, poses significant
challenges but offers a crucial avenue for predicting the impact of potential
interventions. Learned feature representations help us to better infer the
conditional average treatment effect, improving our ability to individually
tailor predictions and target subsets of the general population. In this paper,
we propose and validate a novel representation-based method that uses a
neuroevolutionary approach to remove information from features irrelevant for
predicting the outcome. We demonstrate that this method can yield improved
estimates for heterogeneous treatment effects on standard synthetic examples
and illustrate its use on a real life dataset. We believe that representational
learning is particularly well-suited for removing extraneous information in
causal models and we anticipate future research in this area.

\subsubsection*{Acknowledgements.}

We would like to thank Binjie Lai, Yi-Hong Kuo, and Xiang Wu for their support
and feedback. We are also grateful to the anonymous reviewers for their insights
and suggestions.