Language Models of Code are Few-Shot Commonsense Learners

In plain words

Many tasks ask a language model to produce a graph, such as the steps of a plan and the links between them. The usual trick is to write the graph as a flat text list of pairs or as a Graphviz-style file. Those formats look nothing like the text these models were trained on, so the model often gets the links wrong.

The authors' idea is to write the same graph as a short Python class instead. Each step becomes a method that returns its description, and the links become a list of arrows. A model trained heavily on code is then shown a few finished classes and asked to finish a new one, and the finished code is turned back into a graph that the task's usual scoring can check.

In the paper's tests, this beat a fine-tuned model trained on the task and a text-trained model given the same few examples, on several of the reported metrics. The work covers three structured tasks: event scripts, tracking how entities change through a process, and argument graphs.

TL;DR

• Structured commonsense tasks ask a model to generate a graph (a script, an entity-state trace, or an explanation graph) from text §1.
• The method writes each training pair as a Python class, prompts Codex (code-davinci-002) with a few such classes, and converts its completion back to a graph §2 §2.2.
• Author-reported gains over text-davinci-002 and text-curie-001 with the same prompts, and over a fine-tuned model on several metrics tab. 1 tab. 2 tab. 3 tab. 4.
• The reported numbers use a flatter Python "literal" format, not the node-and-children format the method section describes. The paper and code disagree in several places (see Paper vs code).
• Codex is a closed API; the paper's own Limitations section flags that access may change § Limitations.

Problem

Structured commonsense generation maps a text input to a graph. The paper studies three tasks: script generation (a goal to a graph of steps and their dependencies) §3.2, entity state tracking (actions to a per-entity state trace) §3.3, and argument explanation graphs (belief and argument to a graph of typed edges) §3.4 fig. 1.

The existing approach turns the output graph into text, either as a list of node pairs or as DOT fig. 1. The authors argue that such serialisations differ strongly from pre-training text, so language models find them hard to generate correctly §1. Fine-tuning helps, but it needs task-specific data, and the paper's setting is few-shot: a large model gets a handful of examples and no weight updates §2.2.

Method

Overview. Each training pair (input text T, output graph S) is written as one Python class §2 §2.1. T becomes a class attribute or a comment, and S uses ordinary Python constructs such as objects, attributes and function calls §2.1. A test input is a partial class that holds only the goal. The model completes it, the completion is compiled and read back into a graph, and the dataset's original metrics are computed on that graph §2.2.

Step 1: write examples as code §2.1. The paper's running example is the script "serve the potpies on a plate". Each step becomes a Node object and each edge becomes an entry in the children attribute of its source node fig. 1 §2.1. The authors say this format was found empirically to work best, especially with larger models §2.1. The reported results, however, use a flatter format (see Pitfalls and tab. 15).

Step 2: build the prompt §2.2 §3.1. The prompt is k complete classes, each a (T, S) pair from the training set, concatenated in sequence. For a new input, a partial class with only the goal is appended, and the model writes the rest. The examples are sampled at random from the training set §3.1. The paper says k ranges from 5 to 30 depending on task, because the 4096-token window limits how many fit §3.1.

Step 3: parse the completion §2.2. The authors report that Codex nearly always produced syntactically valid Python, and that the output graphs were similar in structure to the references in nearly all cases §2.2 app. E. The parsed graph is scored with the standard metrics of each dataset.

Task-specific encodings.

• Script generation §3.2 fig. 1. Nodes are Node objects and edges are children lists. Edge prediction gives the nodes in a random order, and the model writes the edges after an # edges comment §3.2.
• Entity state tracking §3.3 fig. 3. The code puts the actions as comments in a main function, defines one state_k variable per participant, defines an init function for the starting states, and defines one function per action that sets each state. The test input is cut off just before init §3.3. The state value UNK marks an unknown location and None marks a non-existent one fig. 3.
• Explanation graphs §3.4. The belief, argument and stance become string variables. The text says each edge is written as an add_edge(source, type, destination) call §3.4, but the reported numbers match the flat string format instead (see Pitfalls and tab. 13). The stance is taken as given from a stance model released with the dataset §3.4.

Dynamic prompts (analysis). The main results use a fixed random prompt. The analysis also tests choosing the k closest training examples for each test input, with a retriever trained by knowledge-similarity tuning (KST) §4 app. B. The loss is given in eq. 1. Results are mixed (see Results).

Edge-prediction scoring §3.2. Precision divides the number of correct predicted edges by the number of predicted edges. Recall divides by the number of reference edges. F1 is the harmonic mean.

# Edge-prediction precision, recall and F1  [§3.2]
# Input:  gold_edges, pred_edges (sets of "a -> b" strings)
# Output: precision, recall, f1 (floats in [0, 1])

def edge_scores(gold_edges, pred_edges):
    # Edges in both sets; duplicates collapse in set()
    # [code: src/eval/edge_pred_eval.py#L55-L56]
    overlap = len(gold_edges & pred_edges)
    # Precision: correct predictions over predicted edges  [§3.2]
    precision = overlap / len(pred_edges) if pred_edges else 0.0
    # Recall: correct predictions over gold edges. The text writes a
    # union here; the code uses the intersection.
    # [§3.2] [code: src/eval/edge_pred_eval.py#L59]
    recall = overlap / len(gold_edges) if gold_edges else 0.0
    # Harmonic mean of precision and recall  [§3.2]
    f1 = 0.0
    if precision + recall > 0:
        f1 = 2 * precision * recall / (precision + recall)
    return precision, recall, f1

KST loss app. B eq. 1. The retriever is a sentence encoder fine-tuned so that cosine similarity between two inputs matches a structural similarity between their target graphs. The paper writes the squared error as eq. 1. The learning rate, batch size and number of epochs are not stated.

L_{\theta} = \big(\cos(f_{\theta}(\mathcal{T}_i), f_{\theta}(\mathcal{T}_j)) - \mathrm{sim}(\mathcal{S}_i, \mathcal{S}_j)\big)^2 eq. 1

# KST training loss for one pair of training examples  [app. B] [eq. 1]
# Input:  text_i, text_j (inputs); struct_sim (target similarity in [0, 1])
# Output: squared error for this pair

def kst_loss(encoder, text_i, text_j, struct_sim):
    # Cosine similarity of the two sentence embeddings  [app. B] [eq. 1]
    emb_i = encoder(text_i)
    emb_j = encoder(text_j)
    cos_ij = cosine(emb_i, emb_j)
    # Squared gap to the structural similarity  [app. B] [eq. 1]
    return (cos_ij - struct_sim) ** 2

Hyperparameters

Values are from the paper's text unless the "Where" column names code. Code values are from the released repository at the commit in the front matter.

Results

Author-reported throughout. Few-shot setups use the number of examples in brackets. "tf" rows are fine-tuned models. Per the paper, the ProScript fine- tuned row is taken from Sakaguchi et al. and not re-run §3.1.

ProScript script generation tab. 1. Reference graph: average degree 1.84, 7.41 nodes, 6.80 edges.

The text says davinci is slightly better on two structural metrics §3.2.

ProScript edge prediction tab. 2.

ProPara entity state tracking tab. 3. Three-shot, averaged over five random prompts tab. 3. Scores are precision, recall and F1.

The text calls the curie result a format failure, since curie "often failed to produce output with the desired format" §3.3. The paper says the method is less than 10 F1 points below the current state of the art, which uses full fine-tuning §3.3 (Ma et al., 2022).

ExplaGraphs argument graphs tab. 4. Evaluated on the validation split. ↓ marks lower-is-better (GED).

Against tf (2500), this method is lower on SECA and EA tab. 4. The text's "across all metrics" claim refers to tf (1500) §3.4.

Where the gain comes from tab. 5. Code prompt versus text prompt, with each model. For ExplaGraphs STCA, codex+text 38.02 and codex+code 45.20. For ProScript edge-prediction F1, codex+text 44.47 and codex+code 56.52. The code prompt helps davinci less (ProScript F1 43.06 to 49.3). The authors conclude that both the model (Codex) and the code format contribute, and that the combination gains more than either alone §4.

Dynamic prompts tab. 6 tab. 8 tab. 9. For ProScript edge prediction, the KST retriever raises F1 from 56.52 to 65.71 tab. 6. For ProScript script generation the gain is marginal (BLEU 25.24 to 25.4) tab. 8. For ExplaGraphs it lowers every metric, because the training set has many near-duplicate examples, and the authors explain it with those duplicates tab. 9 app. B.

Human evaluation tab. 11. Twenty examples, three authors. For ExplaGraphs, graphs from this method were preferred for correctness 38.3% of the time, and davinci's 18.3% (31.7% no preference). For ProScript script generation, correctness was 26.7% against 23.3% (50% no preference) tab. 11. See the Pitfalls on the ExplaGraphs rows.

Model-size and variance checks. Codex-001 is worse than Codex-002 on both tasks with identical prompts tab. 16 tab. 17. Across three seeds, all reported improvements over davinci are said to have p < 0.001 app. I tab. 18 tab. 19 tab. 20 tab. 21.

Estimated model strengths tab. 7. The authors estimate relative model strength from average loss per token on Python code and on Wikipedia text. The paper calls these "educational guesses" about undisclosed models § Limitations app. A.

Paper vs code

The repository is https://github.com/madaan/CoCoGen, linked from the abstract [abstract], which is the basis for treating it as the official code. The commit used here is bd38f89955fcb30eaf4f824191d5be902c725966, the default branch's head at clone time. The repository has no LICENSE file (none in the file listing), so nothing below is quoted verbatim, and code is cited by path and line only.

• Scoring format. The method section and Figure 1 show the node-and-children format for script generation §2.1 fig. 1. The numbers in tab. 1 match the flat "literal" format from the appendix (Literal row in tab. 15), and the literal converter writes one method per step and a get_relations list literal.py. The Tree format is a different converter, proscript-tree get_converter.py tree.py.
• ExplaGraphs format. Section 3.4 describes add_edge(...) calls §3.4. The reported 45.20 STCA matches the Literal row, which stores the whole graph as one string tab. 13. The code for that format is literal.py and the parser is literal.py.
• Edge-prediction format. tab. 14 reports Literal 52.0 and Tree 56.5 F1 for codex-002, but its caption says the literal format performs best. The table does not support the caption. The default edge-prediction job uses the proscript-edge-pred-literal converter README.md edge_pred_literal.py. Which converter gave the tab. 2 numbers is not stated.
• Engine. The paper uses code-davinci-002 §3.1. The README's default command passes --engine text-davinci-002, a natural-language model README.md. Running the README command does not reproduce the paper's Codex setup.
• Decoding. The paper states no decoding settings. The code uses temperature 0.0, top_p 1, zero penalties, best_of 1 and a 280-token limit by default openai_api_wrapper.py query_openai_over_tasks.py.
• Over-length prompts. The paper says prompts are limited by the context window §3.1. The code, on an InvalidRequestError, retries with the first 1 to 6 examples removed query_openai_over_tasks.py. The saved output records the completion and the graph, not how many examples were dropped query_openai_over_tasks.py.
• Prompt sampling. Appendix I says the random seed "decides the order of examples in the prompt" app. I. The shipped sampler uses the seed to pick which examples are sampled make_codex_prompt.py, not their order. Each example ends with # END followed by a blank line make_codex_prompt.py.
• Prompt sizes. The ProScript prompts have 15 examples each, which matches the table prompt.txt prompt.txt. The ExplaGraphs prompt has 32 # END markers against k = 30 in tab. 4 prompt.txt. The ProPara prompt has 5 init functions against 3 in §3.3 prompt.txt.
• Edge-prediction order. The paper says nodes are listed in random order §3.2. The converter reads the order stored in the input edge_pred_literal.py. I did not find where the shuffle happens.
• Edge-prediction metric. The recall formula in the text is a union §3.2. The code uses intersection for both precision and recall edge_pred_eval.py.
• KST encoder and loss. The loss is a mean-squared error between the cosine of two sentence embeddings and a stored similarity, which matches eq. 1 kst_train.py. The label column defaults to sem_sim_score kst_train.py. The encoder default differs from the footnote (see Hyperparameters).
• Converter factory is broken. src/converters/get_converter.py contains unresolved git merge markers at lines 38, 42 and 44 [code: src/converters/get_converter.py#L38-L44]. As checked out, it does not parse. Its imports also use a converters. prefix while the API module uses src. get_converter.py openai_api_wrapper.py.
• ProPara parser. The propara-func converter's python_to_graph is a no-op func.py. I did not trace where the state values are parsed for scoring.
• Dependencies. requirements.txt pins openai == 0.11.0 [code: requirements.txt#L27]. The code calls the older Completions API with an engine argument.

Pitfalls

• Format matters more than the paper suggests. The method section's own example is not the format behind the headline numbers. Implement the literal format for ProScript and ExplaGraphs if you want the reported scores. The paper's appendix shows that the format changes results strongly, for example Literal 25.24 BLEU and Tree 18.96 BLEU for Codex-002 tab. 15.
• Two F1 values for one setting. The ProScript edge-prediction result for this method is F1 56.24 in tab. 2 but 56.52 in tab. 5 and tab. 6. The davinci row also differs: recall 49.30 in tab. 2 and 48.27 in tab. 5. The paper does not reconcile these.
• ProPara numbers differ by section. tab. 3 gives 58.0 and 63.0 F1 for davinci and this method, averaged over five runs. tab. 21 gives 56.9 ±2.4 and 62.8 ±2.4 over three seeds. The text says "at least 3 prompts" §3.1, and the tab. 3 caption says five runs.
• ExplaGraphs seed means differ. The seed-averaged STCA for this method is 44.0 ±1.2 tab. 20, against the 45.20 in tab. 4.
• Human-evaluation rows. The ExplaGraphs relevance row sums to 91.7% and the correctness row to 88.3%, not 100% tab. 11. With 20 examples and three annotators (60 judgments), the rows imply that some judgments are missing.
• Training size. The ProScript fine-tuned baseline is described as using all 3500 samples §3.2. The data statistics list 3,252 training examples app. D.
• Test data access. The ProScript test split came from the authors app. D. A reimplementation cannot score on it without that split.
• Tree and literal disagree on edge prediction. The tab. 14 caption says literal is best, yet Tree scores higher there (56.5 against 52.0 F1).
• Dynamic prompts can hurt. For ExplaGraphs, the KST prompt lowers STCA from 45.2 to 37.47 tab. 9. The duplicates in training data are the authors' explanation app. B. Check for near-duplicates before adding retrieval.
• Codex access and cost. The paper's experiments depend on a closed model and an API whose future access the authors could not guarantee § Limitations. A rerun needs an API key and per-example requests. The code throttles itself to 10 requests per minute by default [code: src/api/query_openai_over_tasks.py#L211].
• English only. The paper evaluates only English datasets § Limitations.
• Figure numbers. Table and figure numbers here follow the order of captions in the LaTeX source, counted by hand. I did not check them against a compiled PDF.

Lineage

• Builds on: Codex (Chen et al., 2021) for the code model §2.2; GPT-3 (Brown et al., 2020) for few-shot prompting and the text baselines §2.2 §3.1; the ProScript, ProPara and ExplaGraphs datasets §3.
• The KST retriever follows the target-similarity tuning of Poesia et al. (2021) app. B.
• Compared against: text-davinci-002 and text-curie-001 with the same prompts §3.1; a fine-tuned model (tf) on each dataset §3.1; the state of the art on ProPara (Ma et al., 2022) §3.3.
• Superseded by: not known from the content of this wiki.