Source-linked AI summary
A Retrieve-and-Edit Framework for Predicting Structured Outputs
Tatsunori B. Hashimoto, Kelvin Guu, Yonatan Oren, Percy Liang
TL;DR
Generating complex outputs can be easier when an existing example is edited, but fixed retrieval metrics and joint retriever-editor training limit this strategy. The paper learns a task-dependent retriever efficiently with a noisy encoder-decoder and combines it with an editor, improving vanilla sequence-to-sequence models on Python autocomplete and Hearthstone. The approach narrows the gap with specialized code-generation models on Hearthstone while substantially improving Python autocomplete over sequence-to-sequence baselines.
Problem
Fixed or generic retrieval may miss task-dependent output similarity, while jointly training the retriever and editor requires an expensive sum over possible retrieved examples.
Method
A noisy encoder-decoder learns embeddings for task-dependent retrieval, after which an editor generates the desired output from the input and retrieved example.
Results
Retrieve-and-edit substantially improves vanilla sequence-to-sequence models on Python autocomplete and Hearthstone, including a 7 BLEU improvement over sequence-to-sequence and retrieval baselines on Hearthstone.
Takeaways & Limitations
The learned retriever combines downstream editability with computationally efficient cosine LSH and narrows the gap between specialized and vanilla code-generation models on Hearthstone.
Takeaways & Limitations
The framework focuses on deterministic retrieval and leaves stochastic and multiple retrievals for future work.
Abstract
from arXiv · showhide
For the task of generating complex outputs such as source code, editing existing outputs can be easier than generating complex outputs from scratch. With this motivation, we propose an approach that first retrieves a training example based on the input (e.g., natural language description) and then edits it to the desired output (e.g., code). Our contribution is a computationally efficient method for learning a retrieval model that embeds the input in a task-dependent way without relying on a hand-crafted metric or incurring the expense of jointly training the retriever with the editor. Our retrieve-and-edit framework can be applied on top of any base model. We show that on a new autocomplete task for GitHub Python code and the Hearthstone cards benchmark, retrieve-and-edit significantly boosts the performance of a vanilla sequence-to-sequence model on both tasks.
1 Introduction
Complex outputs are often variations of previously observed examples, motivating a retrieve-and-edit framework that retrieves a relevant training example and edits it into the desired output. The approach learns task-dependent retrieval efficiently and improves sequence-to-sequence performance on Python autocomplete and Hearthstone code generation.
- Motivation: Complex outputs can be difficult to generate from scratch but may resemble previously observed examples that require only targeted edits.The paper motivates this pattern with natural language generation, music generation, and program synthesis.
- Framework: Retrieve-and-edit first selects a similar training pair and then edits its output prototype into the target output.For example, a retrieved code example can be adapted to a new natural-language description.
- Learned retrieval: Task-specific retrieval is desirable because inputs should be considered similar when their outputs can be transformed through small, easy edits.Fixed hand-crafted or generic retrieval mechanisms may be poorly suited to structured or complex inputs.
- Learned retrieval: Jointly training the retriever and editor would require summing over all possible retrieved inputs for each example, making the straightforward approach prohibitively slow.
- Method: The proposed method trains a noisy encoder-decoder to learn an edit-oriented embedding space, then trains the editor using retrieval under that learned metric.It is computationally efficient and requires no domain knowledge beyond an encoder-decoder with low reconstruction error.
- Results: 14 BLEU points improve a standard sequence-to-sequence model on Python code autocomplete, while Hearthstone performance improves by 7 BLEU and exceeds the best non-AST model by 4 points.The Python dataset contains 76k functions; learned retrieval also improves over a fixed bag-of-words baseline by 6 BLEU.
2 Problem statement
The framework models output generation by retrieving a training example and conditioning an editor on it, while focusing on deterministic retrieval for tractability. Its task-dependent similarity objective is decomposed to avoid summing over all training examples during each optimization step.
- Problem formulation: The task is to learn p_model(y | x), which predicts an output y from an input x drawn from p_data.Examples include generating a 5–15 line code snippet from a natural-language description.
- Generative process: Retrieve-and-edit samples a training pair (x′, y′) with p_ret((x′, y′) | x), then generates y with p_edit(y | x, (x′, y′)).
- Generative process: The overall likelihood combines the editor likelihood p_edit(y | x, (x′, y′)) with the retriever probability p_ret((x′, y′) | x).
- Assumption: The framework focuses on deterministic retrievers that assign all retrieval probability to one training example, leaving stochastic and multiple retrievals for future work.
- Task-dependent similarity: Task-dependent similarity should favor pairs whose retrieved output can be edited into the target with high likelihood, but direct maximum-marginal-likelihood training requires summing over all training examples.That summation is computationally intractable at every training iteration.
- Optimization: The optimization is decomposed by first training the retriever with an oracle editor and then training the editor using the standard maximum-likelihood objective.This avoids the computational difficulty while retaining a task-dependent retriever.
3 Learning to retrieve and edit
The framework learns a task-dependent retriever through a lower-bound objective, then trains an editor using the resulting retriever. A noisy encoder-decoder supplies the embedding metric, enabling efficient nearest-neighbor retrieval without jointly optimizing all retrieved examples.
- Retriever objective: The retriever is trained against an oracle-editor objective because the real editor’s behavior is initially unknown.The oracle editor is the true conditional target distribution given a retrieved example under the retriever–data joint distribution.
- Retriever objective: A variational lower bound separates reconstruction quality from a discrepancy between embeddings of the input and retrieved example.The bound contains a reconstruction term and a KL-divergence discrepancy term.
- Latent-space metric: The method uses a von Mises-Fisher encoder distribution so the KL discrepancy becomes a squared Euclidean distance between unit-sphere embeddings.This choice also provides a bounded KL divergence used to construct a conservative training bound.
- Training procedure: The encoder and decoder are optimized with a reconstruction objective, using reparameterization-gradient minibatch training rather than joint optimization over the retriever.The conservative bound removes dependence on the retriever during encoder-decoder training.
- Training procedure: The final retriever selects the nearest training input under the learned embedding, after which the editor is trained to maximize likelihood given the retrieved example.The implementation uses cosine-LSH nearest-neighbor search, and the editor can use any architecture, including sequence-to-sequence models with attention and copying.
4 Experiments
Experiments evaluate retrieve-and-edit on Python autocomplete and Hearthstone, comparing it with sequence-to-sequence, retrieval, and specialized baselines. Across both tasks, retrieval improves generation performance, while learned task-dependent retrieval outperforms lexical retrieval and produces relevant editable prototypes.
- Python autocomplete: The Python autocomplete dataset contains 76k GitHub functions, with inputs comprising a block comment, function name, arguments, and partial function body.The output is the next token in the function body, evaluated using candidate lists and successive-token completion length.
- Python autocomplete: Retrieve+Edit substantially improves performance over both Seq2Seq and retrieval without editing on Python autocomplete.Table 1 reports candidate-list accuracy and completion-length metrics, with k denoting the number of beam-search candidates.
- Python autocomplete: TaskRetriever significantly outperforms LexicalRetriever on all metrics and is comparable to InputRetriever on BLEU.TaskRetriever is slightly better than InputRetriever on the autocomplete metrics.
- Python autocomplete: In the Python example, Retrieve+Edit uses a relevant retrieved conditional but replaces the unknown condition with a generic conditional.The retrieved prototype supplies structural guidance, whereas Seq2Seq instead defaults to a generic getter function.
- Hearthstone cards benchmark: On Hearthstone, Retrieve+Edit improves BLEU over sequence-to-sequence and retrieval baselines while narrowing the gap to AST-based models.The improvement is 7 BLEU points over the sequence-to-sequence and retrieval baselines and 4 points over the best non-AST method.
- Hearthstone cards benchmark: The Hearthstone prediction differs from the gold standard only by omitting minion_type=MINION_TYPE.none, which is the field’s default setting.The retrieved and desired cards share structural properties, and the editor changes the battlecry modification from ChangeAttack to Silence.
5 Related work
The paper situates retrieve-and-edit among retrieval-based generation, learned embeddings, code generation, and nonparametric or mixture models. Its distinctive setting is conditional generation with deterministic retrieval optimized for editing rather than generic similarity.
- Retrieval models for text generation: Earlier retrieval-based generation systems retrieved and adapted phrases, while later dialogue and translation models trained editors to transform retrieved candidates.
- Retrieval models for text generation: Unlike random-sampling edit-based generation for unconditional tasks, this work addresses conditional sequence generation with a deterministic retriever.
- Embedding models: Noisy autoencoder embeddings, semantic hashing, and cross-modal embeddings learn representations for similarity, but this work uses embeddings to support downstream editability.
- Learned retrieval: Learned retrieval methods based on known correct items or rewards do not directly apply because easy-to-edit retrieved examples are not known in advance.
- Code generation: The framework differs from AST-constrained and copy-based code generation because it uses retrieve-and-edit as a general-purpose performance-boosting method.
- Nonparametric models and mixture models: Retrieve-and-edit can be viewed as nonparametric regression with learned proximity and an editor acting as a learned kernel, or as a mixture model conditioning on retrieved examples.
6 Discussion
The discussion presents learned retrieval with noisy encoder-decoder embeddings as an efficient way to augment standard sequence-to-sequence models for complex outputs. The approach improves Python autocomplete and narrows the gap with specialized models on Hearthstone.
- The framework augments standard sequence-to-sequence models with a learned retriever for generating complex outputs such as source code.
- A noisy encoder-decoder learns retrieval proximity aligned with downstream editability while retaining the computational efficiency of cosine LSH.
- The method narrows the gap between specialized code-generation models and vanilla sequence-to-sequence models on Hearthstone while substantially improving Python autocomplete over sequence-to-sequence baselines.
- Reproducibility: Data and code for reproducing the reported results are available through the CodaLab Worksheets platform.
A KL divergence between two von Mises-Fisher distributions
The section establishes that, for equal-concentration von Mises-Fisher distributions, KL divergence is proportional to squared Euclidean distance between their direction vectors.
- For two vMF distributions with the same concentration parameter κ, their KL divergence is proportional to squared Euclidean distance between direction vectors.
- For directions µ1 and µ2 on the unit sphere S^{d−1}, the proposition relates KL(vMF(µ1, κ), vMF(µ2, κ)) to ∥µ1 − µ2∥2^2.
- The proportionality constant is Cκ, defined through the concentration parameter κ and modified Bessel-function terms.
B.1 Editor
The editor encodes the current input and a retrieved example separately, then uses attention and copying to generate the target sequence. Training also includes an identity-map augmentation for retrieved examples.
- The editor encodes x, x′, and y′ with separate LSTMs and decodes y using their representations plus attention over the inputs.
- Each input component is separately encoded by a 2-layer bidirectional LSTM, and a linear layer combines the final hidden states.
- The editor uses a 4-layer LSTM decoder with attention, while the VAE decoder uses the same decoder architecture without attention.
- A copying mechanism lets the editor copy words from the retrieved example, supporting slot-filling when adapting the retrieved function.
- The dataset is augmented with identity examples mapping (x′, x′, y′) to y′ with probability 0.1.
- Removing the identity augmentation did not substantially change results but increased sensitivity to the number of training epochs.
B.2 Copying
The copying mechanism augments the vocabulary with position-specific copy tokens and combines base-word and copy-token probabilities when words can be copied from the input.
- Copy-token representation: The vocabulary adds 300 copy tokens, with one token corresponding to each possible input position.The number of copy tokens need only reach the largest input length.
- Copy-token representation: Each word embedding concatenates its base-vocabulary vector with the vector for the copy token matching its input position.
- Decoding and training: The decoder applies a softmax over the full vocabulary, including copy tokens, during both training and testing.
- Decoding and training: During training, a target word’s probability sums the base-word probability and the matching copy-token probability when that word appears in the input.The same additive definition applies to the target word’s log-likelihood.
C Github dataset construction
The GitHub dataset was constructed by querying BigQuery for Python files meeting minimum-size and documentation-style criteria, with the query capped at five million files.
- Selection criteria: The dataset retrieval targeted Python files of at least 100 bytes containing a function with a reStructuredText-style docstring.
- Query construction: The BigQuery query joined GitHub repository file metadata with repository contents using matching file and content identifiers.
- Selection criteria: The regex required a triple-quoted content span containing “returns:”, and the query limited results to 5,000,000 files.