PaperlayerView exact revision on arXiv

Source-linked AI summary

CodeIt: Self-Improving Language Models with Prioritized Hindsight Replay

Natasha Butt, Blazej Manczak, Auke Wiggers, Corrado Rainone, David W. Zhang, Michaël Defferrard, Taco Cohen

arXiv:2402.04858v2cs.AIcs.CLcs.LG

TL;DR

ARC remains difficult because existing systems struggle with sparse rewards and inter-task generalization from few-shot examples. CodeIt addresses this by alternating program sampling and hindsight relabeling with prioritized experience replay, solving 59/400 ARC evaluation tasks and achieving state-of-the-art performance. The method also supports solution refinement, with shorter programs found for 53% of solved tasks.

  • Problem

    ARC requires generalization from few-shot examples, but existing methods struggle with sparse rewards and fail to use experience for inter-task generalization.

    Source: §1 · pp. 1–2

  • Method

    CodeIt alternates program sampling with hindsight relabeling and learning from prioritized experience replay to train a language model to write ARC programs.

    Source: §1 · p. 2, §2 · p. 2, §2.2 · p. 3

  • Results

    59/400 ARC evaluation tasks are solved, achieving state-of-the-art performance against symbolic and neural baselines.

    Source: §6 · p. 9

  • Takeaways & Limitations

    CodeIt demonstrates self-improvement through program-space reasoning, inter-task generalization, and solution refinement over repeated iterations.

    Source: §1 · p. 2, §6 · p. 9

    Direct comparisons with competition results are difficult because those results use hidden test sets and may condition the DSL on evaluation tasks.

    Source: §4.1 · p. 7

Abstract

from arXiv · show

Large language models are increasingly solving tasks that are commonly believed to require human-level reasoning ability. However, these models still perform very poorly on benchmarks of general intelligence such as the Abstraction and Reasoning Corpus (ARC). In this paper, we approach ARC as a programming-by-examples problem, and introduce a novel and scalable method for language model self-improvement called Code Iteration (CodeIt). Our method iterates between 1) program sampling and hindsight relabeling, and 2) learning from prioritized experience replay. By relabeling the goal of an episode (i.e., the target program output given input) to the realized output produced by the sampled program, our method effectively deals with the extreme sparsity of rewards in program synthesis. Applying CodeIt to the ARC dataset, we demonstrate that prioritized hindsight replay, along with pre-training and data-augmentation, leads to successful inter-task generalization. CodeIt is the first neuro-symbolic approach that scales to the full ARC evaluation dataset. Our method solves 15% of ARC evaluation tasks, achieving state-of-the-art performance and outperforming existing neural and symbolic baselines. Our code is available at https://github.com/Qualcomm-AI-research/codeit .

1. Introduction

ARC tests generalization from few-shot grid examples, yet existing neural and symbolic methods struggle to use experience across tasks. CodeIt addresses this gap by combining program-based expert iteration with hindsight relabeling and prioritized replay, achieving 59/400 evaluation tasks.

1. Introduction

  • State-of-the-art GPT-4-based neural approaches solve only 12% of ARC evaluation tasks, while humans solve 80% of a random subset in user studies.ARC is designed to rely on innate core knowledge systems rather than knowledge memorized from broad pretraining.Source: §1 · p. 1
  • ARC requires agents to infer abstractions from demonstration grids and generalize both to test examples and across tasks.Its tasks use few-shot input-output examples involving objects, actions, number, and space.Source: §1 · p. 1
  • Existing neural and symbolic approaches generally fail to use experience for inter-task generalization, motivating expert iteration over programs rather than output grids.Expert iteration alternates between gathering data with exploration and improving the policy on new experiences.Source: §1 · p. 2
  • Sparse rewards make ARC program synthesis sample-inefficient, while hindsight relabeling creates artificial trajectories but changing data distributions risk catastrophic forgetting.These limitations motivate combining hindsight relabeling with prioritized experience replay.Source: §1 · p. 2
  • CodeIt alternates program sampling with hindsight relabeling and learning with prioritized experience replay, generating data without human intervention.Unlike sampling-and-filtering approaches, it learns from all program samples and prioritizes experiences that solve real tasks.Source: §1 · p. 2
  • 59/400 ARC evaluation tasks are solved by CodeIt, achieving state-of-the-art performance through abstraction-based experience and inter-task generalization.The paper also reports that CodeIt analyzes discovered programs and performs solution refinement over time.Source: §1 · p. 2

2. Method

CodeIt treats ARC as programming by examples and alternates program sampling with hindsight relabeling and learning from prioritized replay. Its design combines a DSL-constrained pretrained language model with augmented training data and replay experiences that remain useful even when sampled programs do not solve the original task.

2. Method

  • ARC is formulated as programming by examples: a policy generates DSL programs from demonstration input-output grids to match their mappings.The search set contains tasks for which the method seeks corresponding programs.Source: §2 · p. 2, §2.2 · p. 3

2.1. Design choices

  • The method uses Hodel’s DSL, restricted to a language designed from the ARC training split rather than evaluation-set priors.The DSL provides grid manipulation, filling, and pixel-location functions.Source: §2.1 · pp. 2–3
  • A pretrained 220-million-parameter CodeT5+ encoder-decoder receives demonstration examples and autoregressively generates the corresponding program.Demonstrations may be truncated to fit the encoder context window.Source: §2.1 · p. 3
  • The grids are represented as object-centric text that lists colored cells by coordinates, reducing tokens for sparse grids without using a custom tokenizer.This representation is intended to remain human-interpretable and is especially beneficial for larger or sparse grids.Source: §2.1 · p. 3, Figure 3 · p. 3

2.2. The Code Iteration Algorithm

  • The policy is initialized on ground-truth ARC programs, then repeatedly trained using the original, augmented, and replay-buffer experiences.Each full sampling-and-learning pass is called a meta-iteration.Source: §2.2 · p. 3
  • Randomly mutated programs augment the initial training data, teaching DSL syntax and demonstration interpretation before iterative sampling begins.This augmentation is intended to improve the quality of early policy samples.Source: §2.2 · p. 3
  • CodeIt alternates sampling programs with hindsight relabeling and learning from prioritized replay, using valid realized program outputs as training experiences.Relabeling ensures syntactically correct programs contribute experiences even when their outputs differ from the target outputs.Source: §2.2 · pp. 3–4
  • Prioritized replay samples experiences more often when their program outputs match real demonstration outputs, while the single continually updated policy has no formal improvement guarantee.The authors report that continual updates did not cause practical problems, but acknowledge that later iterations can in principle perform worse.Source: §2.2 · pp. 3–4

3. Experiments

The experiments evaluate CodeIt on ARC using controlled program-sampling and learning procedures, compare it with custom and literature baselines, and ablate its main components. CodeIt achieves state-of-the-art evaluation performance, while prioritized replay, pretraining, demonstrations, and mutation-based augmentation materially affect learning.

3.1. Setup

  • The setup initializes CodeIt from ARC training programs, 19,200 synthetically mutated programs, and a 400-task evaluation search set, then samples up to 24 valid programs per task per meta-iteration.Learning samples 10,000 experiences from the train and augmented sets and 90,000 from the replay buffer at each meta-iteration.Source: §3.1 · p. 4

3.2. Main results on ARC eval set

  • CodeIt eventually surpasses mutation baselines as cumulative performance rises beyond solutions already present in the augmented training set.Figure 4 compares cumulative performance against custom and literature baselines as the number of sampled programs increases.Source: §3.2 · pp. 4–5, Figure 4 · p. 5
  • CodeIt outperforms prior symbolic and neural baselines on the full ARC evaluation set after 100 meta-iterations, achieving state-of-the-art performance.The evaluation uses pass@3 by default and also reports pass@1 on ARC Eval 412 for comparison with Gendron et al. (2023).Source: §3.2 · p. 5, Table 1 · p. 5
  • CodeIt solves 29 tasks jointly with the mutation baseline, 23 additional tasks alone versus 13 alone for the mutation baseline, and finds shorter programs on shared solutions.It also uses different primitives on tasks solved by both methods.Source: §3.2 · p. 5

3.3. Ablations

  • Removing policy feedback reduces cumulative performance, while removing hindsight relabeling causes stagnation after about 30 meta-iterations compared with CodeIt.The no-ExIt ablation still exceeds mutation-only learning, showing that supervised learning from mutation experiences provides some inter-task generalization.Source: §3.3 · p. 6
  • Removing prioritized replay sharply reduces policy performance, while removing pretraining, multiple demonstrations, or mutation augmentation slows or weakens learning.Uniform replay increases forgetting; one-demo input strongly reduces both performance measures, and no mutation produces slower but continuing improvement.Source: §3.3 · pp. 6–7

4. Related work

Prior ARC methods either operate on subsets, solve tasks independently, or use self-improvement approaches that do not adequately address ARC’s sparse rewards and inter-task generalization. CodeIt builds on expert iteration, hindsight replay, and language-model-generated synthetic experience to address these gaps.

4.1. Abstraction and Reasoning Corpus (ARC)

  • Full-set ARC methods solve tasks independently, with reported results including 36/400, 27/400, and 49/412 evaluation tasks.These approaches generally rely on custom DSL search or pretrained language models, while some subset methods do not scale to the full evaluation set.Source: §4.1 · p. 7
  • Competition results are difficult to compare because some participants condition DSLs on the public evaluation set, and competition leaderboards use hidden test sets.The paper therefore distinguishes competition performance from results evaluated on the ARC evaluation set.Source: §4.1 · p. 7

4.2. Expert Iteration

  • Expert iteration alternates policy-guided experience gathering with imitation learning, commonly using computationally intensive search experts.CodeIt instead samples programs from a language-model policy and learns from program experiences in a programming-by-examples setting.Source: §4.2 · p. 7
  • Unlike DreamCoder’s correctness filtering and growing DSL, CodeIt stores distilled knowledge in the model and uses hindsight relabeling with prioritized replay.This design targets inter-task generalization while avoiding reliance on a computationally intensive search procedure.Source: §4.2 · p. 7

4.3. Self Improving Large Language Models

  • Self-improving language-model methods have succeeded across synthetic-data domains, and this work extends that strategy to the challenging ARC domain.The related approaches include theorem proving, question answering, mathematical reasoning, translation, and code generation.Source: §4.3 · pp. 7–8

5. Discussion

The discussion attributes CodeIt’s ARC performance to combining prior knowledge with experience, especially hindsight relabeling, prioritized replay, and scalable policy components. It also identifies hand-designed DSL and evaluation infrastructure as important scope limitations.

5. Discussion

  • 59/400 unseen ARC tasks are solved, and the discussion explains this performance through expert iteration that combines prior knowledge with learned experience.ARC’s limited data, complex tasks, and varying demonstrations and input dimensionality make this combination particularly relevant.Source: §5 · p. 8
  • Hindsight relabeling creates an implicit curriculum from easy programs to more complex ones, while prioritized sampling reduces catastrophic forgetting.The replay buffer contains many relabeled programs, including experiences farther from the target tasks.Source: §5 · p. 8
  • CodeIt for ARC relies on a hand-designed DSL, an interpreter for automatic evaluation, and an initial set of ground-truth programs.The authors identify approaches that learn tabula rasa or form their own DSL as an important area for future research.Source: §5 · p. 8
  • A neuro-symbolic ablation outperforms the symbolic mutation baseline, while CodeIt performs better still, indicating contributions from DSL knowledge, learning, and expert iteration.Training remains possible without the pretrained LLM and mutation procedure, but at a slower pace.Source: §5 · p. 8
  • Unlike prior approaches that use priors without new experience, CodeIt combines pretrained language models, likelihood-based training, interpreters, and experience to support scalable inter-task generalization.The learned ARC policy also permits using a smaller language model than approaches relying on GPT-4.Source: §5 · p. 8

6. Conclusion

CodeIt is presented as a scalable self-improvement method using prioritized hindsight replay, achieving state-of-the-art ARC performance while learning and refining programs across tasks.

6. Conclusion

  • 59/400 ARC evaluation tasks are solved, achieving state-of-the-art performance against symbolic and neural baselines.The conclusion identifies scalability and learning from experience as key components for success on ARC.Source: §6 · p. 9
  • Hindsight relabeling improves sample efficiency by 40%, while prioritized training ameliorates catastrophic forgetting.These components address the sparse-reward setting by retaining and emphasizing useful program experiences.Source: §6 · p. 9
  • Later iterations find shorter programs for 53% of solved tasks, showing that CodeIt can refine solutions over time.The conclusion characterizes this behavior as reasoning in program space and generalizing between tasks.Source: §6 · p. 9

Impact Statement

The paper frames CodeIt as iterative program learning from interpreter feedback, while documenting both its societal risks and the concrete algorithmic stages used to initialize, sample, relabel, and train the policy.

Impact Statement

  • The approach trains a language-model policy to write programs for programming-by-example problems using iterative interpreter feedback.Its intended setting is data-scarce programming-by-example learning.Source: Impact Statement · p. 9
  • Small datasets may amplify biases in language-model training data, and code-generating models could automate malicious code.These are the paper’s stated potential negative societal impacts.Source: Impact Statement · p. 9

A.1. CodeIt Algorithm

  • CodeIt initializes with augmented training data, an empty optimal-program set, a policy, and a replay buffer before repeated meta-iterations.The algorithm requires training and search sets, a policy, a replay buffer, and an optimal-program set.Source: §A.1 · p. 12
  • Each sampling stage generates programs for search tasks, adds hindsight-relabeled input-output examples to replay, and records programs that solve tasks.The subsequent learning stage samples from replay and augmented and original training data, then trains the policy for one epoch.Source: §A.1 · p. 12

A.2. Program and Task Mutation

The method evolves ARC training tasks by mutating programs and their demonstrations, while evaluating solutions under standard ARC pass@3 or pass@1 settings.

A.2. Program and Task Mutation

  • Programs are mutated by selecting a random line and replacing a function, an argument, or a function call according to typed alternatives.The mutation baseline evaluates each mutated program on tasks in the search set.Source: §A.2 · p. 12, §A.3 · p. 13

A.3. Task Representation

  • The representation compresses grids to reduce token usage, but truncates large demonstrations and programs to fit encoder and decoder context windows.Demonstration inputs and outputs share the encoder budget, while programs use the decoder context window.Source: §A.3 · pp. 12, 14
  • Mutated programs execute on each demonstration input, and valid output grids are added as new tasks to the evolving population.The procedure initializes mutated demonstrations, executes the mutated program, validates all outputs, and appends the resulting task.Source: §A.3 · p. 13

A.4. ARC evaluation

  • The default ARC Eval Set contains 400 tasks and reports pass@3, selecting three demonstration-performing programs and solving a task only when all test outputs are correct.Ties favor shorter programs; an alternative ARC Eval 412 setting creates one task per test example and reports pass@1.Source: §A.4 · p. 14

B. Experiment details

Experiments use an internal validation split and a training-derived ARC DSL, with resource and initialization choices designed to tune CodeIt without using the evaluation split.

B.1. Resources

  • CodeIt experiments ran for up to 120 hours on an NVIDIA A100 80GB.Source: §B.1 · p. 14

B.2. Hyperparameter tuning

  • The ARC training data is split into 311 training and 89 validation examples with roughly matched program-length difficulty.The validation set is used for hyperparameter tuning and as the sampling search set, avoiding tuning directly on ARC evaluation.Source: §B.2 · p. 14

B.3. Hyperparamaters chosen on internal validation set

  • The method’s hyperparameters were optimized on the custom validation set before application to ARC evaluation.Source: §B.3 · p. 14

B.4. Domain Specific Language

  • The adopted DSL was designed from the training set without evaluation-set inspection, but some evaluation programs may remain unsolvable with its available primitives.It includes grid transformations, counting, and operations on foreground objects.Source: §B.4 · p. 15
  • CodeIt training is initialized with 400 hand-designed solution programs covering the ARC training tasks, including programs as long as 58 lines.Source: §B.4 · p. 15

C. ConceptARC

On the easier, concept-grouped ConceptARC benchmark, CodeIt’s strengths and weaknesses vary by task concept relative to GPT-4 baselines.

C. ConceptARC

  • CodeIt outperforms GPT-4 on extend to boundary, extract objects, filled not filled, and same different concepts.The comparison uses greedy and sampled pass@3 evaluation against GPT-4 zero-shot and few-shot-refine baselines.Source: C. ConceptARC · p. 15, Figure 8 · p. 16
  • CodeIt substantially underperforms GPT-4 on order, clean up, and count concepts.Source: C. ConceptARC · p. 15
  • CodeIt performs best on object-interaction tasks and worst on numerical or logic-based tasks, despite the DSL containing ordering and counting primitives.Source: C. ConceptARC · p. 15

D. Scaling Study

Scaling improves CodeIt performance with larger CodeT5 models, while analyses show that solution length and primitive frequency constrain what the system learns.

D. Scaling Study

  • Performance increases significantly from 60M to 220M parameters and more modestly from 220M to 770M parameters.The 770M model was evaluated for only 85 meta-iterations because of computation constraints.Source: D. Scaling Study · p. 16, Figure 9 · p. 17
  • LoRA finetuning with Mistral-7B produced fewer than 9% solved tasks after 200k sampled programs, below CodeT5 220M at close to 11% and 770M at around 12%.The authors suggest larger models may overfit the small dataset and leave broader LoRA and backbone searches for future work.Source: D. Scaling Study · p. 16

E.1. Failure Cases on Custom Validation Set

  • CodeIt almost always solves validation tasks whose ground-truth programs contain 10 or fewer lines, but does not solve medium-length tasks with 11–15 lines.For longer tasks, it often learns the primitives without finding a complete solution.Source: §E.1 · p. 16, Figure 10 · p. 17, Figure 11 · p. 17

E.2. DSL Learning

  • Some ARC evaluation primitives are learned faster than others, with object-returning primitives learned sooner than numeric-returning primitives.Primitives are often not learned when they occur in few training programs or in long programs.Source: §E.2 · p. 18, Figure 12 · p. 18

E.3. CodeIt compared with mutation baselines

CodeIt solves more ARC evaluation tasks than the mutation baseline and generally finds shorter, less redundant programs. Across meta-iterations, it can refine solutions by finding shorter programs later, although its prioritization does not explicitly favor program length.

E.3. CodeIt compared with mutation baselines

  • CodeIt solves 52/400 ARC evaluation tasks, including 23 not solved by the mutation baseline.The comparison uses the best-performing CodeIt run against the mutation d1 baseline.Source: §E.3 · p. 18
  • Among 29 tasks solved by both methods, shortest programs differed in 24 cases; CodeIt produced a longer program in 1 case, making its program shorter in the other 23.The Mutation baseline often includes redundant lines, while CodeIt programs appear more concise, less complex, and shorter overall.Source: §E.3 · pp. 18–19
  • CodeIt programs are generally shorter than mutation-baseline programs, with no obvious correlation between task-representation length and program length.The task representations may be truncated beyond the 1024-token encoder context, yet some truncated representations still suffice for solving tasks.Source: §E.3 · p. 18, Figure 13 · p. 19, Figure 14 · p. 19

E.4. CodeIt over time

  • Across 45 tasks, CodeIt finds shorter solutions over time in 24 tasks, despite its priority not explicitly incorporating program length.After finding a solution, later meta-iterations can produce both longer and shorter solutions, leaving room for program refinement.Source: §E.4 · p. 19, Figure 15 · p. 19, Table 6 · p. 21
Loading 2402.04858v2…