[Paper Review] CORAL: COde RepresentAtion Learning with Weakly-Supervised Transformers for Analyzing Data Analysis

Ge Zhang, Mike A. Merrill+3 more|arXiv (Cornell University)|Aug 28, 2020

Software Engineering Research4 citations

TL;DR

CORAL proposes a weakly supervised transformer-based model that jointly encodes code from abstract syntax trees and surrounding natural language comments to classify Jupyter notebook cells into data analysis stages (e.g., import, wrangling, exploration, modeling, evaluation). It achieves a 38% accuracy gain over expert heuristics and enables large-scale analysis of 118,000 notebooks, revealing that notebook composition correlates with publication citation impact.

ABSTRACT

Large scale analysis of source code, and in particular scientific source code, holds the promise of better understanding the data science process, identifying analytical best practices, and providing insights to the builders of scientific toolkits. However, large corpora have remained unanalyzed in depth, as descriptive labels are absent and require expert domain knowledge to generate. We propose a novel weakly supervised transformer-based architecture for computing joint representations of code from both abstract syntax trees and surrounding natural language comments. We then evaluate the model on a new classification task for labeling computational notebook cells as stages in the data analysis process from data import to wrangling, exploration, modeling, and evaluation. We show that our model, leveraging only easily-available weak supervision, achieves a 38% increase in accuracy over expert-supplied heuristics and outperforms a suite of baselines. Our model enables us to examine a set of 118,000 Jupyter Notebooks to uncover common data analysis patterns. Focusing on notebooks with relationships to academic articles, we conduct the largest ever study of scientific code and find that notebook composition correlates with the citation count of corresponding papers.

Motivation & Objective

• To enable large-scale, in-depth analysis of scientific code without requiring expensive expert annotations.
• To develop a model that classifies code snippets into stages of the data analysis pipeline using only weak supervision from simple heuristics.
• To uncover patterns in scientific data analysis across millions of Jupyter notebooks and relate them to publication impact.
• To create a publicly available dataset of 1,840 annotated notebook cells for evaluation of code stage classification.

Proposed method

• CORAL uses a novel masked attention mechanism to jointly model natural language comments and structured code from abstract syntax trees.
• It employs a weakly supervised training objective based on five simple heuristics to replace costly manual labeling.
• The model combines weak supervision with unsupervised topic modeling in a multi-task optimization framework to improve representation learning.
• It leverages pre-trained transformer architectures to learn joint code and comment representations for stage classification.
• A graph neural network component processes code structure, integrating syntax tree information with contextual comments.
• The model is fine-tuned on a new dataset of 1,840 expert-annotated cells from 100 notebooks, used solely for evaluation.

Experimental results

Research questions

• RQ1 Do scientific notebooks follow consistent sequential patterns in their data analysis workflows?
• RQ2 How does the distribution of code stages in notebooks relate to the citation impact of associated academic papers?
• RQ3 Can weak supervision from simple heuristics enable high-accuracy classification of code stages without expert-annotated labels?
• RQ4 Are there domain-specific differences in how notebook content correlates with publication impact?

Key findings

• CORAL achieves a 38% improvement in accuracy over expert-supplied heuristics on classifying code into data analysis stages.
• Notebooks linked to academic papers receive 21.88 times more citations on average than those without linked notebooks.
• A one-standard-deviation increase in stage entropy (uniformity of stage distribution) is associated with a 2.11-fold increase in expected citations.
• In computer science and mathematics, a higher proportion of data wrangling code is linked to lower citation counts, while this effect is absent in biological sciences.
• Notebooks with a high volume of exploration-stage code are associated with 35% fewer citations, suggesting such notebooks may reflect incomplete or exploratory work.
• The model enables large-scale analysis of 118,000 Jupyter notebooks, revealing that most cells transition to a different stage, indicating a dynamic, non-linear data analysis process.