Using Small Language Models to Reverse-Engineer Machine Learning Pipelines Structures

TL;DR

This paper investigates the use of Small Language Models to automatically identify and understand the structure of machine learning pipelines from source code, aiming to improve flexibility and reliability over existing methods.

Contribution

It demonstrates that Small Language Models can effectively classify ML pipeline components, offering a scalable and adaptable approach compared to manual labeling and traditional classifiers.

Findings

SLMs outperform traditional classifiers in pipeline structure classification

Performance varies with taxonomy definitions, affecting classification accuracy

SLMs provide insights aligning with data science practices from prior studies

Abstract

Background: Extracting the stages that structure Machine Learning (ML) pipelines from source code is key for gaining a deeper understanding of data science practices. However, the diversity caused by the constant evolution of the ML ecosystem (e.g., algorithms, libraries, datasets) makes this task challenging. Existing approaches either depend on non-scalable, manual labeling, or on ML classifiers that do not properly support the diversity of the domain. These limitations highlight the need for more flexible and reliable solutions. Objective: We evaluate whether Small Language Models (SLMs) can leverage their code understanding and classification abilities to address these limitations, and subsequently how they can advance our understanding of data science practices. Method: We conduct a confirmatory study based on two reference works selected for their relevance regarding current state-of-the-art's limitations. First, we compare several SLMs using Cochran's Q test. The best-performing model is then evaluated against the reference studies using two distinct McNemar's tests. We further analyze how variations in taxonomy definitions affect performance through an additional Cochran's Q test. Finally, a goodness-of-fit analysis is conducted using Pearson's chi-squared tests to compare our insights on data science practices with those from prior studies.