Structured Code Representations Enable Data-Efficient Adaptation of Code Language Models

TL;DR

This paper demonstrates that representing code as parse trees and further pre-training on these structures significantly improves data efficiency and performance of pre-trained code models, especially with limited training data.

Contribution

The work introduces a method to adapt pre-trained code models using program structures like parse trees, enhancing data efficiency without changing the model architecture.

Findings

Improved performance on code tasks with structural pre-training.

Significant gains in low-data scenarios.

Structural adaptation benefits models pre-trained only on plain text.

Abstract

Current language models tailored for code tasks often adopt the pre-training-then-fine-tuning paradigm from natural language processing, modeling source code as plain text. This approach, however, overlooks the unambiguous structures inherent in programming languages. In this work, we explore data-efficient adaptation of pre-trained code models by further pre-training and fine-tuning them with program structures. Specifically, we represent programs as parse trees -- also known as concrete syntax trees (CSTs) -- and adapt pre-trained models on serialized CSTs. Although the models that we adapt have been pre-trained only on the surface form of programs, we find that a small amount of continual pre-training and fine-tuning on CSTs without changing the model architecture yields improvements over the baseline approach across various code tasks. The improvements are found to be particularly significant when there are limited training examples, demonstrating the effectiveness of integrating program structures with plain-text representation even when working with backbone models that have not been pre-trained with structures.