Structure-aware Fine-tuning for Code Pre-trained Models

TL;DR

This paper tackles the challenge of transferring task-specific structural knowledge to CodePTMs during fine-tuning. It introduces Structure-aware Fine-tuning (SAT), which uses a structure loss based on the Sinkhorn divergence between the Transformer’s attention matrix and an AST-derived distance matrix, combined with multi-task fine-tuning and a learnable structure encoder; the final objective is $ abla\mathcal{L}_{\text{final}}=\mathcal{L}_{\text{task}}+\alpha\mathcal{L}_{\text{structure}}$ applied to downstream code-generation tasks. Empirical results across four CodePTMs and two generation tasks show consistent improvements, with larger gains in low-resource settings, and CAT-probing/case studies corroborate that SAT enhances the model’s ability to capture code structure. The method is plug-and-play and architecture-agnostic, offering a practical route to boost structural knowledge absorption without modifying model internals. The work has implications for improving program understanding and generation in real-world scenarios where labeled data may be scarce.

Abstract

Over the past few years, we have witnessed remarkable advancements in Code Pre-trained Models (CodePTMs). These models achieved excellent representation capabilities by designing structure-based pre-training tasks for code. However, how to enhance the absorption of structural knowledge when fine-tuning CodePTMs still remains a significant challenge. To fill this gap, in this paper, we present Structure-aware Fine-tuning (SAT), a novel structure-enhanced and plug-and-play fine-tuning method for CodePTMs. We first propose a structure loss to quantify the difference between the information learned by CodePTMs and the knowledge extracted from code structure. Specifically, we use the attention scores extracted from Transformer layer as the learned structural information, and the shortest path length between leaves in abstract syntax trees as the structural knowledge. Subsequently, multi-task learning is introduced to improve the performance of fine-tuning. Experiments conducted on four pre-trained models and two generation tasks demonstrate the effectiveness of our proposed method as a plug-and-play solution. Furthermore, we observed that SAT can benefit CodePTMs more with limited training data.

Figures (7)

• Figure 1: Compared to regular fine-tuning, CodePTMs can effectively capture task-specific code structure knowledge through structure-aware fine-tuning by structure loss.
• Figure 2: Schematic illustration of generation of the attention matrix and the distance matrix. (a) Generation of the attention matrix. The code sequence is input into the multi-head self-attention block of CodePTMs, then the scaled dot-product values are computed to generate attention matrix. (b) Generation of the distance matrix. We first parse the raw code into AST, and calculate the shortest path between leaves on the AST to generate distance matrix.
• Figure 3: The architecture of the SAT method. Given a code snippet, we parsed it into an AST, and the distance matrix is obtained by computing the shortest distances between leaf nodes. The structural information is extracted using the struct encoder. The code snippet is then input into the CodePTMs, from which the attention matrix is extracted from the multi-head attention blocks. The sinkhorn divergence between the distance matrix and the attention matrix is calculated to obtain the structure loss.
• Figure 4: Average results of SAT on code summarization and code translation tasks.
• Figure 5: Performance improvement of fine-tuning with SAT on limited training datasets.