Analysing the Behaviour of Tree-Based Neural Networks in Regression Tasks

TL;DR

This paper investigates how tree-based neural networks perform in regression tasks like predicting code execution time, introduces a dual-transformer model leveraging source code and ASTs, and demonstrates its superior performance over existing models.

Contribution

The paper extends tree-based neural network models to regression tasks, proposes a novel dual-transformer architecture with cross-attention, and evaluates its effectiveness on real-world datasets.

Findings

Dual-transformer outperforms existing tree-based models

Models show limitations in regression tasks

Dual-transformer demonstrates robustness across datasets

Abstract

The landscape of deep learning has vastly expanded the frontiers of source code analysis, particularly through the utilization of structural representations such as Abstract Syntax Trees (ASTs). While these methodologies have demonstrated effectiveness in classification tasks, their efficacy in regression applications, such as execution time prediction from source code, remains underexplored. This paper endeavours to decode the behaviour of tree-based neural network models in the context of such regression challenges. We extend the application of established models--tree-based Convolutional Neural Networks (CNNs), Code2Vec, and Transformer-based methods--to predict the execution time of source code by parsing it to an AST. Our comparative analysis reveals that while these models are benchmarks in code representation, they exhibit limitations when tasked with regression. To address these deficiencies, we propose a novel dual-transformer approach that operates on both source code tokens and AST representations, employing cross-attention mechanisms to enhance interpretability between the two domains. Furthermore, we explore the adaptation of Graph Neural Networks (GNNs) to this tree-based problem, theorizing the inherent compatibility due to the graphical nature of ASTs. Empirical evaluations on real-world datasets showcase that our dual-transformer model outperforms all other tree-based neural networks and the GNN-based models. Moreover, our proposed dual transformer demonstrates remarkable adaptability and robust performance across diverse datasets.