GEPCode: A Context-Aware 1M-Parameters Graph-Based Language Model for Source Code
Keywords: graph neural network, graph language model, source code optimization
Abstract: The pursuit of optimal conditions for software execution poses a complex challenge. This task can be automated by harnessing the structured nature of programming languages, especially from compiler intermediate representations of code (IR). The manipulation of source code using Large Language Models (LLMs) is a thriving area of study in Natural Language Processing (NLP) literature. However, in this study we illustrate how we can circumvent the need for exceedingly large models by employing domain-specific language models. These models have a reduced number of parameters but retain the ability to capture the relationships within source code elements. We introduce GEPCode, a graph neural network designed to model IR with the flexibility to adapt to new tasks. This flexibility is obtained through special "meta" nodes, that allow for the representation of additional task-dependent contextual information. Pre-training is performed by solving node and graph-level tasks, resulting in a general language model. After a fine-tuning phase on two downstream tasks, Device Mapping and Algorithm Classification, we achieve average accuracy results of 88.9% (NVIDIA) and 92.3% (AMD) for the former and 97.2% for the latter. Comparing our methodology with state-of-the-art models trained from scratch, our results are similar or better, yet providing a more flexible model. Moreover, we achieve similar accuracy results in downstream tasks compared to state-of-the-art pre-trained language models based on Transformers, while utilizing 100 times fewer parameters.
Primary Area: learning on graphs and other geometries & topologies
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 11356
Loading