Preprint
Machine Learning

CodeGen2

May 1, 2023

0

Citations

0

Influential Citations

Venue

2023

Year

Abstract

Proposes an approach to make the training of LLMs for program synthesis more efficient by unifying key components of model architectures, learning methods, infill sampling, and data distributions

Analysis

Why This Paper Matters

CodeGen2 tackles a core challenge in training large language models (LLMs) for program synthesis: how to efficiently combine the strengths of different architectures, learning objectives, and data sources. As AI practitioners increasingly deploy models that must both generate and understand code, the paper’s systematic evaluation of design choices offers actionable insights. It challenges assumptions like the "free lunch" of infill sampling and the necessity of separate encoder-decoder architectures, providing a clearer path toward unified models.

How It Works

The paper evaluates four hypotheses through controlled experiments:

  • Model Architecture: It compares causal decoders (uni-directional attention) with PrefixLM (bi-directional attention on a prefix, uni-directional on the rest). The goal is to unify encoder and decoder capabilities.

  • Learning Algorithm: A mixture of causal language modeling (CLM) and span corruption is used. Span corruption masks random spans of tokens, and the model must reconstruct them, combining next-token prediction with denoising.

  • Sampling Procedure: The authors test whether training with infill objectives (e.g., PSM: prefix, suffix, middle) degrades left-to-right generation. They hypothesize a "free lunch" where infill adds no cost.

  • Data Distributions: Models are trained on a 50/50 mix of natural language (The Pile) and code (The Stack) to see if multi-domain training harms single-domain performance.

Figure 1

Key technical details:

  • PrefixLM splits input into a prefix (bi-directional attention) and context (causal attention).
  • Span corruption uses a dynamic ratio of tokens to mask, with file-level corruption to avoid document boundary issues.
  • No task tokens are prefixed; the choice between CLM and span corruption is random (50% probability).

Results

Comparison of causal decoder and Prefix-LM on various tasks.

  • Architecture: PrefixLM performs competitively with causal decoders on most tasks but shows worse results on the Stack dataset, suggesting a negative effect from masked loss. It also struggles to outperform smaller encoder-only models on discriminative tasks.
  • Infill: Training with infill objectives reduces HumanEval pass@1 by about 1 point, disproving the "free lunch" hypothesis.
  • Objective: A simple mixture of CLM and span corruption works well, but using UL2 (which treats a large portion as prefix) degrades HumanEval performance. Modifying denoiser hyperparameters recovers performance, emphasizing the importance of maximizing gradient-update tokens.
  • Data Mixing: The mixed model (Mix) does not outperform single-domain models given the same compute budget but matches them, showing substantial improvement over single-domain models evaluated on the other domain.

Results on LAMBADA and PIQA for NL and HumanEval for PL.

  • Multi-Epoch Training: CodeGen2.5 (7B, trained 5 epochs on StarCoderData) shows significant pass@k improvements over CodeGen2-7B, though the contribution of individual factors (more tokens, data augmentation, extended learning rate) is unclear.

Pass@k on HumanEval under various temperatures.

Significance

CodeGen2 provides practical lessons for training LLMs on code and natural language. The key takeaway is that simple, carefully chosen objectives (CLM + span corruption) can be effective, while architectural innovations like PrefixLM may not justify their complexity. The paper also highlights trade-offs: infill sampling degrades left-to-right generation, and multi-domain training requires careful resource allocation. For practitioners, this means focusing on objective design and data mixing strategies rather than chasing architectural novelties. The multi-epoch training results suggest that repeated exposure to data, combined with span corruption as augmentation, can yield substantial gains, though the exact mechanism needs further study.