PaperlayerView exact revision on arXiv

Source-linked AI summary

CodeGen2: Lessons for Training LLMs on Programming and Natural Languages

Erik Nijkamp, Hiroaki Hayashi, Caiming Xiong, Silvio Savarese, Yingbo Zhou

arXiv:2305.02309v2cs.LG

TL;DR

LLM training for programming and natural-language tasks is costly because practitioners must choose among architectures, objectives, sampling procedures, and data distributions. This paper evaluates a unified recipe across these components using 1B-parameter experiments. It finds mixed results, including promising data mixing and strong multi-epoch-training evidence, but no measured Prefix-LM benefit and no free lunch for infill.

  • Problem

    LLM training is costly and leaves practitioners without obvious choices among architectures, learning algorithms, and data distributions.

  • Method

    The study performs extensive 1B-parameter ablations to unify architecture, objectives, sampling procedures, and data distributions.

  • Results

    The study finds no measurable Prefix-LM improvement, no free lunch for infill, promising natural/programming-language mixing, and strong multi-epoch-training evidence.

  • Takeaways & Limitations

    The exploration yields five lessons, a final training recipe, and competitive open-source models and implementation for practitioners.

  • Takeaways & Limitations

    The authors speculate that insufficient gradient updates for a target language may negatively affect Prefix-LM, but the threshold or progression remains unknown.

Abstract

from arXiv · show

Large language models (LLMs) have demonstrated remarkable abilities in representation learning for program synthesis and understanding tasks. The quality of the learned representations appears to be dictated by the neural scaling laws as a function of the number of model parameters and observations, while imposing upper bounds on the model performance by the amount of available data and compute, which is costly. In this study, we attempt to render the training of LLMs for program synthesis more efficient by unifying four key components: (1) model architectures, (2) learning methods, (3) infill sampling, and, (4) data distributions. Specifically, for the model architecture, we attempt to unify encoder and decoder-based models into a single prefix-LM. For learning methods, (i) causal language modeling, (ii) span corruption, (iii) infilling are unified into a simple learning algorithm. For infill sampling, we explore the claim of a "free lunch" hypothesis. For data distributions, the effect of a mixture distribution and multi-epoch training of programming and natural languages on model performance is explored. We conduct a comprehensive series of empirical experiments on 1B LLMs, for which failures and successes of this exploration are distilled into five lessons. We will provide a final recipe for training and release CodeGen2 models in size 1B, 3.7B, 7B, and, 16B parameters, along with the training framework as open-source: https://github.com/salesforce/CodeGen.

1 INTRODUCTION

The paper targets the cost and complexity of choosing separate architectures, objectives, sampling procedures, and data distributions for programming and natural-language tasks. It evaluates a unified recipe and finds mixed results: some components succeed, while Prefix-LM and free-lunch infilling lack supporting evidence.

  • 1.1 MOTIVATION: COST OF LLMS: LLM training requires costly choices among attention masks, learning algorithms, task capabilities, and data distributions.These unresolved choices motivate seeking a single universal model and recipe.
  • 1.2 GOALS: REDUCE COST BY UNIFICATION AND OPEN-SOURCE: The proposed recipe unifies architecture, learning objective, left-to-right and infill sampling, and data distributions into one model.The stated goal is competitive performance across synthesis and understanding tasks.
  • 1.3 FINDINGS: A MIXED BAG OF RESULTS: The study evaluates unification hypotheses through extensive ablations on 1B-parameter language models.The hypotheses cover architecture, learning algorithm, sampling procedure, and data distributions.
  • 1.3 FINDINGS: A MIXED BAG OF RESULTS: Prefix-LM showed no quantified benefit over a causal-decoder baseline, while the simple objective mixture maintained zero-shot performance.The evidence is limited to the study’s evaluation tasks.
  • 1.3 FINDINGS: A MIXED BAG OF RESULTS: The free-lunch hypothesis for infill sampling was unsupported, whereas mixing natural and programming languages showed promising evidence.The study also reports strong results for multi-epoch training.
  • 1.4 CONTRIBUTIONS: LESSONS, RECIPE, AND OPEN-SOURCE: The work contributes five lessons, a simple mixture objective, an open-source implementation, and planned open-source CodeGen2 models.The objective mixes uncorrupted and within-file span-corruption sequences with next-token prediction.

2 METHOD: FROM UNIFICATION TO ABLATION

The method defines a common evaluation framework for synthesis, infill, in-context learning, and understanding, then compares modeling choices across attention, objectives, sampling, and data. It contrasts bidirectional encoder representations with causal decoder representations while targeting efficient task coverage under a shared recipe.

  • 2 METHOD: FROM UNIFICATION TO ABLATION: The study seeks to unify learning methods and model architectures while maintaining task performance under an identical compute budget.Its evaluation spans program synthesis, understanding, and in-context learning.
  • 2 METHOD: FROM UNIFICATION TO ABLATION: Left-to-right synthesis conditions on a function signature and docstring, with HumanEval measuring synthesized-program quality.The program is sampled autoregressively from the prompt.
  • 2 METHOD: FROM UNIFICATION TO ABLATION: Infill synthesis conditions on tokens before and after a missing span, evaluated with HumanEval-Infill.The model samples the intervening body from both left and right context.
  • 2 METHOD: FROM UNIFICATION TO ABLATION: Few-shot in-context learning predicts a label from example input-label pairs, using XSum for evaluation.The setup supplies n examples before predicting the next label.
  • 2 METHOD: FROM UNIFICATION TO ABLATION: Understanding evaluation uses bidirectional representations for defect detection and related tasks on CodeXGLUE and SuperGLUE.Bidirectional representations allow each input token to use all input tokens simultaneously.
  • 2 METHOD: FROM UNIFICATION TO ABLATION: The architectural comparison contrasts encoder attention to all sequence positions with decoder attention restricted to previous positions.The learning-algorithm discussion links encoder models to span corruption and decoder models to next-token prediction.
  • 2 METHOD: FROM UNIFICATION TO ABLATION: The sampling and data components address infill by rearranging sequence parts and study how maximum-likelihood training models data distributions.Infill is motivated by editing code within existing files, where future context can be useful.

3 RESULTS: LESSONS AND RECIPE

This section introduces the empirical results and conclusions of the unification and ablation study. Its findings are distilled into five lessons.

  • 3 RESULTS: LESSONS AND RECIPE: The empirical results and conclusions are distilled into five lessons from the unification and ablation study.

3.1 LESSON 1: PREFIX-LM’S BENEFIT IS QUESTIONABLE

Prefix-LM combines bidirectional encoder attention with causal decoding, but its benefits depend on the data, representation task, and objective. It matches causal decoding on Python-only training yet underperforms on mixed-language data and does not replace smaller encoder-only models.

  • Architecture: Prefix-LM uses bidirectional attention for an initial non-causal sequence part and a causal mask for the remainder.This architecture is intended to combine encoder-style representations with autoregressive decoding.
  • Data: HumanEval scores were generally on-par with a causal decoder during BigPython training, except during the first 100,000 steps.The comparison used models trained for the same number of steps, with non-causal prefixes ranging from 0 to 0.9 of sequence length.
  • Data: 2 points worse pass@k resulted when the same experiment used the Stack, where Python comprised 9.4% of tokens.The Prefix-LM was strictly worse throughout pretraining on this mixed-language dataset.
  • Representation: Prefix-LM generally outperformed a same-objective causal decoder on SuperGLUE but did not match much smaller CodeBERT or RoBERTa-large models.The 1B model therefore did not provide sufficiently informative representations to justify replacing smaller encoder-only models.
  • Objective: Simplifying UL2 denoiser settings improved HumanEval consistently over the original hyperparameters.The S-denoiser share increased from 14% to 50%, while its average prefix length decreased from 75% to 50%.

3.2 LESSON 2: INFILL IS NOT A FREE LUNCH

The paper tests whether infill capability can be added without sacrificing left-to-right generation. Its causal decoder with mixed CLM and PSM infilling lost about 1 HumanEval pass@1 point, so the free-lunch hypothesis was not supported.

  • Motivation: Infill supports editing tokens within a file by conditioning on both preceding and following context.The paper contrasts this with left-to-right sampling, which uses only previous tokens.
  • Method: The experiment trained a causal decoder with a mixture of causal language modeling and PSM sequence-reordering infilling.PSM denotes Prefix, Suffix, Middle sequence reordering, following Bavarian et al.’s experimental setting.
  • Results: About 1 point decrease in HumanEval pass@1 occurred relative to a causal decoder trained only with causal language modeling.The observation agreed with the reported finding that infill is not free and motivated leaving careful reimplementation for future work.

3.3 LESSON 3: OBJECTIVE CAN BE SIMPLE, YET NEEDS TO BE CAREFULLY CHOSEN

The paper favors a simple mixed objective for supporting both left-to-right and infill generation, while emphasizing that objective design must be matched carefully to the target task. CodeGen2 applies this recipe across four model sizes on permissively filtered Stack data.

  • Lesson: The authors conclude that Prefix-LM is not ideal because of poor multilingual performance and sub-par representation utility, while infill remains close to non-infill performance.These findings motivate using a causal decoder recipe rather than relying on full architectural unification.
  • Recipe: The recipe aligns sequence formats, removes task-token complexity, and assumes uniform task-mixture and span-length distributions.Span corruption is selected as the base infill objective, following InCoder, with a different span-selection strategy.
  • Recipe: A 0.5 probability selects span corruption or causal language modeling for each sample in the mixed objective.The method omits UL2 task tokens because their presence produced no notable differences.
  • Evaluation: CodeGen2 evaluates the recipe at 1B, 3.7B, 7B, and 16B parameters using permissively filtered Stack v1.1 data.Evaluation covers HumanEval and HumanEval-Infill, with end-of-mask truncation used for the infill measurement.

3.4 LESSON 4: MULTI-MODAL DATA MIXING

Mixing natural and programming-language data improves cross-domain performance relative to domain-mismatched training, but does not outperform domain-matched baselines under the same compute budget.

  • Motivation: Programming-language models can perform poorly on natural-language tasks despite implicit modality mixing in comments, documents, and code snippets.The study motivates adding a sufficient natural-language learning signal through explicit data mixing.
  • Evaluation: A mixed-data causal decoder is evaluated zero-shot on HumanEval, LAMBADA, and PIQA across training steps.The evaluation reports HumanEval pass@1, LAMBADA accuracy and perplexity, and PIQA accuracy against single-domain baselines.
  • Findings: Mix does not outperform other baselines for the evaluated tasks given the same compute budget.Each domain receives less exposure under the mixture, reducing performance relative to domain-matched models such as PL on HumanEval.
  • Findings: Mix performs closely to domain-matched models while substantially improving over domain-mismatched baselines from early training.The reported pattern supports mixing modalities when compute is constrained and the model serves both domains.
  • Implication: A simple mixture of training data is considered promising for efficient learning across both natural and programming-language domains.The authors suggest that slightly longer mixed training could yield one competitive model for both domains.

3.5 LESSON 5: MULTI-EPOCH TRAINING

The study tests whether repeated observations can improve learning when finite data limits scaling. A 7B CodeGen2.5 model trained for five epochs shows significant pass@k gains over the 400B-token CodeGen2-7B baseline, though the source of the gain remains unresolved.

  • Motivation: Repeated observations are hypothesized to improve learning when finite data limits one-epoch pre-training.Span corruption may act as data augmentation, while an extended learning-rate decay supports higher learning rates for longer.
  • Method: A 7B CodeGen2.5 model is trained on repeated Star-CoderData for 1.4 trillion tokens, equivalent to five epochs.The learning-rate schedule decays over the total training steps.
  • Results: Significant gains in HumanEval pass@k are observed over the baseline CodeGen2-7B model trained for 400B tokens.The improvement continues despite no exposure to new data after 280B tokens.
  • Open question: The contributions of additional tokens, span-corruption augmentation, and prolonged higher learning rates remain unresolved.The authors identify these factors as requiring additional ablations in future work.

4 CONCLUSION: LESSONS AND OPEN-SOURCE

The paper consolidates empirical successes and failures from attempts to unify architecture, objectives, sampling, and data distributions for efficient LLM training. It concludes with a practitioner-oriented recipe and plans to open-source the implementation and CodeGen2 model family.

  • Lessons: The exploration distills positive and negative findings across architecture, objectives, sampling procedures, and data distributions into five lessons.The authors state that satisfactory unification was not achieved, but the findings and recipe may still help practitioners.
  • Lessons: The reported lessons include limited Prefix-LM improvements, no free lunch for infill sampling, promising data mixing, and strong multi-epoch-training evidence.They also identify a simple mixture of causal language modeling and within-file span corruption as sufficient.
  • Open-source: The authors plan to open-source the training implementation and CodeGen2 models at 1B, 3.7B, 7B, and 16B parameters.The release is intended to facilitate future research and community use.
Loading 2305.02309v2…